update parent cookie from todo

Author: thuyen

lua/orgmode/files/headline.lua

@@ -357,12 +357,14 @@ end
 function Headline:set_todo(keyword)
   local todo, node = self:get_todo()
   if todo then
-    return self:_set_node_text(node, keyword)
+    self:_set_node_text(node, keyword)
+    return self:update_parent_cookie_from_todo()
   end
 
   local stars = self:_get_child_node('stars')
   local _, level = stars:end_()
-  return self:_set_node_text(stars, ('%s %s'):format(('*'):rep(level), keyword))
+  self:_set_node_text(stars, ('%s %s'):format(('*'):rep(level), keyword))
+  return self:update_parent_cookie_from_todo()
 end
 
 memoize('get_todo')
@@ -904,28 +906,47 @@ function Headline:set_cookie(num, denum)
   return self
 end
 
-function Headline:update_cookie()
-  local section = self:node():parent()
-  if not section then
+function Headline:update_parent_cookie_from_todo()
+  local parent = self:get_parent_headline()
+  if not parent or not parent.headline 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
+  -- Update the cookie
+  return parent:update_cookie()
+end
+
+function Headline:update_cookie()
+  local num, denum = 0, 0
+
+  -- Count checked boxes from all lists
+  local section = self:node():parent()
+  if section then
+    local body = section:field('body')[1]
+    if body then
+      for node in body:iter_children() do
+        if node:type() == 'list' then
+          local boxes = self:child_checkboxes(node)
+          denum = denum + #boxes
+          local checked_boxes = vim.tbl_filter(function(box)
+            return box:match('%[%w%]')
+          end, boxes)
+          num = num + #checked_boxes
+        end
+      end
     end
   end
 
+  -- Count done children headlines
+  local peers = self:get_child_headlines()
+  local dones = vim.tbl_filter(function(h)
+    return h:is_done()
+  end, peers)
+  num = num + #dones
+  denum = denum + #peers
+
   -- Update the cookie
-  return self:set_cookie(num_checked_boxes, num_boxes)
+  return self:set_cookie(num, denum)
 end
 
 function Headline:child_checkboxes(list_node)