Update cookies from multiple list

Author: thuyen

lua/orgmode/files/elements/listitem.lua

@@ -71,7 +71,7 @@ function Listitem:update_checkbox(action)
   else
     local parent_headline = self.file:get_closest_headline_or_nil()
     if parent_headline then
-      parent_headline:update_cookie(parent_list)
+      parent_headline:update_cookie()
     end
   end
 end

lua/orgmode/files/headline.lua

@@ -890,22 +890,42 @@ function Headline:get_cookie()
   return self:_parse_title_part('%[%d?%d?%d?%%%]')
 end
 
-function Headline:update_cookie(list_node)
-  local total_boxes = self:child_checkboxes(list_node)
-  local checked_boxes = vim.tbl_filter(function(box)
-    return box:match('%[%w%]')
-  end, total_boxes)
-
+function Headline:set_cookie(num, denum)
   local cookie = self:get_cookie()
   if cookie then
     local new_cookie_val
     if self.file:get_node_text(cookie):find('%%') then
-      new_cookie_val = ('[%d%%]'):format((#checked_boxes / #total_boxes) * 100)
+      new_cookie_val = ('[%d%%]'):format((num / denum) * 100)
     else
-      new_cookie_val = ('[%d/%d]'):format(#checked_boxes, #total_boxes)
+      new_cookie_val = ('[%d/%d]'):format(num, denum)
     end
     return self:_set_node_text(cookie, new_cookie_val)
   end
+  return self
+end
+
+function Headline:update_cookie()
+  local section = self:node():parent()
+  if not section then
+    return self
+  end
+
+  -- Go through all the lists in this headline and gather checked_boxes
+  local num_boxes, num_checked_boxes = 0, 0
+  local body = section:field('body')[1]
+  for node in body:iter_children() do
+    if node:type() == 'list' then
+      local boxes = self:child_checkboxes(node)
+      num_boxes = num_boxes + #boxes
+      local checked_boxes = vim.tbl_filter(function(box)
+        return box:match('%[%w%]')
+      end, boxes)
+      num_checked_boxes = num_checked_boxes + #checked_boxes
+    end
+  end
+
+  -- Update the cookie
+  return self:set_cookie(num_checked_boxes, num_boxes)
 end
 
 function Headline:child_checkboxes(list_node)