From b8101e919cd1f4f67e9066eb81b6b43b03ead0dd Mon Sep 17 00:00:00 2001
From: oncomouse <oncomouse@gmail.com>
Date: Fri, 9 Jun 2023 09:39:27 -0500
Subject: [PATCH 1/3] feat: update_cookie() now works for both list and
 headline children

---
 lua/orgmode/org/mappings.lua        |  1 +
 lua/orgmode/treesitter/headline.lua | 68 +++++++++++++++++++++++------
 lua/orgmode/treesitter/listitem.lua |  2 +-
 3 files changed, 57 insertions(+), 14 deletions(-)

diff --git a/lua/orgmode/org/mappings.lua b/lua/orgmode/org/mappings.lua
index a2dc0e4e5..7e79b44d1 100644
--- a/lua/orgmode/org/mappings.lua
+++ b/lua/orgmode/org/mappings.lua
@@ -564,6 +564,7 @@ function OrgMappings:handle_return(suffix)
     local content = config:respect_blank_before_new_entry({ string.rep('*', item.level) .. ' ' .. suffix })
     vim.fn.append(linenr, content)
     vim.fn.cursor(linenr + #content, 0)
+    require('orgmode.treesitter.headline'):new(item.node):update_cookie()
     return vim.cmd([[startinsert!]])
   end
 
diff --git a/lua/orgmode/treesitter/headline.lua b/lua/orgmode/treesitter/headline.lua
index 77a94baac..1b8d49b20 100644
--- a/lua/orgmode/treesitter/headline.lua
+++ b/lua/orgmode/treesitter/headline.lua
@@ -205,12 +205,15 @@ function Headline:set_todo(keyword)
   local current_todo = self:todo()
   if current_todo then
     tree_utils.set_node_text(current_todo, keyword)
+    self:update_cookie()
     return
   end
 
   local stars = self:stars()
   local text = ts.get_node_text(stars, 0)
   tree_utils.set_node_text(stars, string.format('%s %s', text, keyword))
+  self:refresh()
+  self:update_cookie()
 end
 
 function Headline:item()
@@ -436,22 +439,61 @@ function Headline:cookie()
   return self:parse('%[%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)
-
-  local cookie = self:cookie()
-  if cookie then
-    local new_cookie_val
-    if ts.get_node_text(cookie, 0):find('%%') then
-      new_cookie_val = ('[%d%%]'):format((#checked_boxes / #total_boxes) * 100)
+function Headline:update_cookie()
+  local total = 0
+  local done = 0
+  local target = self
+
+  -- Determine the target (headline with the cookie). This could be parent
+  -- headline, as Headline:set_todo() will likely be called on headlines
+  -- whose parents have cookies.
+  local cookie = target:cookie()
+  if not cookie then
+    local parent_section = target.headline:parent():parent()
+    if parent_section:child(0):type() == 'headline' then
+      local parent_headline = Headline:new(parent_section:child(0))
+      cookie = parent_headline:cookie()
+      if cookie ~= nil then
+        target = parent_headline
+      else
+        return nil
+      end
     else
-      new_cookie_val = ('[%d/%d]'):format(#checked_boxes, #total_boxes)
+      return nil
+    end
+  end
+
+  -- Parse the children of the headline's parent section for child headlines and lists:
+  for _, node in pairs(ts_utils.get_named_children(tree_utils.find_parent_type(self.headline, 'section'))) do
+    -- The child is a list:
+    if node:type() == 'body' and node:child(0):type() == 'list' then
+      local total_boxes = target:child_checkboxes(node:child(0))
+      local checked_boxes = vim.tbl_filter(function(box)
+        return box:match('%[%w%]')
+      end, total_boxes)
+      total = total + #total_boxes
+      done = done + #checked_boxes
+    end
+    -- The child is a section:
+    if node:type() == 'section' and node:child(0):type() == 'headline' then
+      local hl = Headline:new(node:child(0))
+      local _, word, is_done = hl:todo()
+      if word ~= nil then
+        total = total + 1
+      end
+      if is_done then
+        done = done + 1
+      end
     end
-    tree_utils.set_node_text(cookie, new_cookie_val)
   end
+
+  local new_cookie_val
+  if ts.get_node_text(cookie, 0):find('%%') then
+    new_cookie_val = ('[%d%%]'):format((total == 0 and 0 or done / total) * 100)
+  else
+    new_cookie_val = ('[%d/%d]'):format(done, total)
+  end
+  tree_utils.set_node_text(cookie, new_cookie_val)
 end
 
 function Headline:child_checkboxes(list_node)
diff --git a/lua/orgmode/treesitter/listitem.lua b/lua/orgmode/treesitter/listitem.lua
index ed7ead036..b43a1f8f7 100644
--- a/lua/orgmode/treesitter/listitem.lua
+++ b/lua/orgmode/treesitter/listitem.lua
@@ -69,7 +69,7 @@ function Listitem:update_checkbox(action)
   else
     local parent_headline = tree_utils.closest_headline()
     if parent_headline then
-      Headline:new(parent_headline):update_cookie(parent_list)
+      Headline:new(parent_headline):update_cookie()
     end
   end
 end

From a8d19342a3f35be9ac33d5dc8cf4993924ce308c Mon Sep 17 00:00:00 2001
From: oncomouse <oncomouse@gmail.com>
Date: Sun, 18 Jun 2023 09:50:05 -0500
Subject: [PATCH 2/3] fix: move is_done check inside of word ~= nil check

---
 lua/orgmode/treesitter/headline.lua | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/lua/orgmode/treesitter/headline.lua b/lua/orgmode/treesitter/headline.lua
index 1b8d49b20..390921997 100644
--- a/lua/orgmode/treesitter/headline.lua
+++ b/lua/orgmode/treesitter/headline.lua
@@ -480,9 +480,9 @@ function Headline:update_cookie()
       local _, word, is_done = hl:todo()
       if word ~= nil then
         total = total + 1
-      end
-      if is_done then
-        done = done + 1
+        if is_done then
+          done = done + 1
+        end
       end
     end
   end

From 69749c14a815a35b75f9d0c855edda195282e971 Mon Sep 17 00:00:00 2001
From: oncomouse <oncomouse@gmail.com>
Date: Sun, 18 Jun 2023 10:16:07 -0500
Subject: [PATCH 3/3] fix: problem with parent section calculation when cookie
 is in a parent headline

---
 lua/orgmode/treesitter/headline.lua | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/lua/orgmode/treesitter/headline.lua b/lua/orgmode/treesitter/headline.lua
index 390921997..4b0eb53a3 100644
--- a/lua/orgmode/treesitter/headline.lua
+++ b/lua/orgmode/treesitter/headline.lua
@@ -447,9 +447,11 @@ function Headline:update_cookie()
   -- Determine the target (headline with the cookie). This could be parent
   -- headline, as Headline:set_todo() will likely be called on headlines
   -- whose parents have cookies.
+  local parent_section = tree_utils.find_parent_type(target.headline, 'section')
   local cookie = target:cookie()
   if not cookie then
-    local parent_section = target.headline:parent():parent()
+    -- We need to check the next section up:
+    parent_section = target.headline:parent():parent()
     if parent_section:child(0):type() == 'headline' then
       local parent_headline = Headline:new(parent_section:child(0))
       cookie = parent_headline:cookie()
@@ -464,7 +466,7 @@ function Headline:update_cookie()
   end
 
   -- Parse the children of the headline's parent section for child headlines and lists:
-  for _, node in pairs(ts_utils.get_named_children(tree_utils.find_parent_type(self.headline, 'section'))) do
+  for _, node in pairs(ts_utils.get_named_children(parent_section)) do
     -- The child is a list:
     if node:type() == 'body' and node:child(0):type() == 'list' then
       local total_boxes = target:child_checkboxes(node:child(0))