From 43fd594eb12379cde3a96f8b224c2ae3506aa042 Mon Sep 17 00:00:00 2001
From: Thuyen Ngo <vanthuyen@gmail.com>
Date: Sat, 26 Oct 2024 09:50:30 -0700
Subject: [PATCH 1/2] Update cookies from multiple list

---
 lua/orgmode/files/elements/listitem.lua |  2 +-
 lua/orgmode/files/headline.lua          | 30 +++++++++++++++++++------
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/lua/orgmode/files/elements/listitem.lua b/lua/orgmode/files/elements/listitem.lua
index ec09e9358..b57a107a3 100644
--- a/lua/orgmode/files/elements/listitem.lua
+++ b/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
diff --git a/lua/orgmode/files/headline.lua b/lua/orgmode/files/headline.lua
index 1bde5b19b..7ca98a59a 100644
--- a/lua/orgmode/files/headline.lua
+++ b/lua/orgmode/files/headline.lua
@@ -890,22 +890,38 @@ 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: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
   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_checked_boxes / num_boxes) * 100)
     else
-      new_cookie_val = ('[%d/%d]'):format(#checked_boxes, #total_boxes)
+      new_cookie_val = ('[%d/%d]'):format(num_checked_boxes, num_boxes)
     end
     return self:_set_node_text(cookie, new_cookie_val)
   end
+  return self
 end
 
 function Headline:child_checkboxes(list_node)

From 675257c95372b735a93f95f10c35cedbd0ee66b3 Mon Sep 17 00:00:00 2001
From: Thuyen Ngo <vanthuyen@gmail.com>
Date: Sun, 9 Mar 2025 09:13:57 -0700
Subject: [PATCH 2/2] add a test

---
 tests/plenary/ui/mappings/checkbox_spec.lua | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/tests/plenary/ui/mappings/checkbox_spec.lua b/tests/plenary/ui/mappings/checkbox_spec.lua
index 03df64ac9..69240511b 100644
--- a/tests/plenary/ui/mappings/checkbox_spec.lua
+++ b/tests/plenary/ui/mappings/checkbox_spec.lua
@@ -151,4 +151,25 @@ describe('Checkbox mappings', function()
       '- [ ] checkbox item',
     }, vim.api.nvim_buf_get_lines(0, 0, 6, false))
   end)
+
+  it('should update headline cookies from multiple lists', function()
+    helpers.create_file({
+      '* Test orgmode [/]',
+      'First List',
+      '- [ ] checkbox item',
+      '- [ ] checkbox item',
+      'Second List',
+      '- [ ] checkbox item',
+    })
+    vim.fn.cursor(4, 1)
+    vim.cmd([[exe "norm \<C-space>"]])
+    assert.are.same({
+      '* Test orgmode [1/3]',
+      'First List',
+      '- [ ] checkbox item',
+      '- [X] checkbox item',
+      'Second List',
+      '- [ ] checkbox item',
+    }, vim.api.nvim_buf_get_lines(0, 0, 6, false))
+  end)
 end)