From 6d3a853110b6f3f9b36ba43119e30ca2571bae09 Mon Sep 17 00:00:00 2001
From: Tray Dennis <tman1300@aol.com>
Date: Mon, 9 May 2022 20:02:55 -0400
Subject: [PATCH 1/6] Add updating of checkbox stat

---
 lua/orgmode/org/mappings.lua        |  4 ++++
 lua/orgmode/treesitter/headline.lua | 18 ++++++++++++++++++
 lua/orgmode/treesitter/list.lua     | 21 +++++++++++++++++++++
 lua/orgmode/utils/treesitter.lua    | 22 +++++++++++++++++++++-
 tests/plenary/ui/mappings_spec.lua  | 11 +++++++++++
 5 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 lua/orgmode/treesitter/list.lua

diff --git a/lua/orgmode/org/mappings.lua b/lua/orgmode/org/mappings.lua
index 63eb4d11d..5507f7093 100644
--- a/lua/orgmode/org/mappings.lua
+++ b/lua/orgmode/org/mappings.lua
@@ -12,6 +12,7 @@ local ts_utils = require('nvim-treesitter.ts_utils')
 local utils = require('orgmode.utils')
 local tree_utils = require('orgmode.utils.treesitter')
 local Headline = require('orgmode.treesitter.headline')
+local List = require('orgmode.treesitter.list')
 
 ---@class OrgMappings
 ---@field capture Capture
@@ -162,6 +163,9 @@ function OrgMappings:toggle_checkbox()
   checkbox = checkbox:gsub('%[[%sXx%-]?%]$', new_val)
   local new_line = line:gsub(pattern, checkbox)
   vim.fn.setline('.', new_line)
+  local headline = Headline:new(tree_utils.closest_headline())
+  local list = List:new(tree_utils.closest_list())
+  headline:update_checkbox_status(list)
 end
 
 function OrgMappings:timestamp_up_day()
diff --git a/lua/orgmode/treesitter/headline.lua b/lua/orgmode/treesitter/headline.lua
index e6c90bc54..b0cd9f7a9 100644
--- a/lua/orgmode/treesitter/headline.lua
+++ b/lua/orgmode/treesitter/headline.lua
@@ -120,6 +120,24 @@ function Headline:remove_closed_date()
   tree_utils.set_node_text(dates['CLOSED'], '', true)
 end
 
+function Headline:checkbox_status()
+  return self:parse('%[%d?/%d?%]')
+end
+
+function Headline:update_checkbox_status(list)
+  local checkbox_status = self:checkbox_status()
+  if not checkbox_status then
+    return
+  end
+
+  local checkboxes = list:checkboxes()
+  local checked_boxes = vim.tbl_filter(function(box)
+    return box:match('%[%w%]')
+  end, checkboxes)
+  local new_status = ('[%d/%d]'):format(#checked_boxes, #checkboxes)
+  tree_utils.set_node_text(checkbox_status, new_status)
+end
+
 -- @return tsnode, string
 function Headline:parse(pattern)
   local match = ''
diff --git a/lua/orgmode/treesitter/list.lua b/lua/orgmode/treesitter/list.lua
new file mode 100644
index 000000000..6fcc761ac
--- /dev/null
+++ b/lua/orgmode/treesitter/list.lua
@@ -0,0 +1,21 @@
+local ts_utils = require('nvim-treesitter.ts_utils')
+local tree_utils = require('orgmode.utils.treesitter')
+local query = vim.treesitter.query
+
+local List = {}
+
+function List:new(list_node)
+  local data = { list = list_node }
+  setmetatable(data, self)
+  self.__index = self
+  return data
+end
+
+function List:checkboxes()
+  return vim.tbl_map(function(node)
+    local text = query.get_node_text(node, 0)
+    return text:match('%[.%]')
+  end, ts_utils.get_named_children(self.list))
+end
+
+return List
diff --git a/lua/orgmode/utils/treesitter.lua b/lua/orgmode/utils/treesitter.lua
index 8857d980a..8fdb08cff 100644
--- a/lua/orgmode/utils/treesitter.lua
+++ b/lua/orgmode/utils/treesitter.lua
@@ -2,6 +2,11 @@ local ts_utils = require('nvim-treesitter.ts_utils')
 local config = require('orgmode.config')
 local M = {}
 
+function M.current_node()
+  local window = vim.api.nvim_get_current_win()
+  return ts_utils.get_node_at_cursor(window)
+end
+
 -- walks the tree to find a headline
 function M.find_headline(node)
   if node:type() == 'headline' then
@@ -19,7 +24,22 @@ end
 -- returns the nearest headline
 function M.closest_headline()
   vim.treesitter.get_parser(0, 'org'):parse()
-  return M.find_headline(ts_utils.get_node_at_cursor(vim.api.nvim_get_current_win()))
+  return M.find_headline(M.current_node())
+end
+
+function M.find_list(node)
+  if node:type() == 'list' then
+    return node
+  end
+  if node:type() == 'body' then
+    return nil
+  end
+  return M.find_list(node:parent())
+end
+
+function M.closest_list()
+  vim.treesitter.get_parser(0, 'org'):parse()
+  return M.find_list(M.current_node())
 end
 
 -- @param front_trim boolean
diff --git a/tests/plenary/ui/mappings_spec.lua b/tests/plenary/ui/mappings_spec.lua
index b65bc6b94..e6fb3086c 100644
--- a/tests/plenary/ui/mappings_spec.lua
+++ b/tests/plenary/ui/mappings_spec.lua
@@ -1596,4 +1596,15 @@ describe('Mappings', function()
       '*** to be refiled',
     }, vim.api.nvim_buf_get_lines(0, 0, 5, false))
   end)
+
+  it('should update the checklist status on a headline', function()
+    helpers.load_file_content({
+      '* Test orgmode [/]',
+      '- [ ] checkbox item',
+      '- [ ] checkbox item',
+    })
+    vim.fn.cursor(2, 1)
+    vim.cmd([[exe "norm \<C-space>"]])
+    assert.are.same('* Test orgmode [1/2]', vim.fn.getline(1))
+  end)
 end)

From e09227a263bdb15b6feb0e565a2040aeb5b0e278 Mon Sep 17 00:00:00 2001
From: Tray Dennis <tman1300@aol.com>
Date: Mon, 9 May 2022 20:07:53 -0400
Subject: [PATCH 2/6] update terminology to "cookie" (from org manual)

---
 lua/orgmode/org/mappings.lua        | 2 +-
 lua/orgmode/treesitter/headline.lua | 6 +++---
 tests/plenary/ui/mappings_spec.lua  | 2 +-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/lua/orgmode/org/mappings.lua b/lua/orgmode/org/mappings.lua
index 5507f7093..529b3f98d 100644
--- a/lua/orgmode/org/mappings.lua
+++ b/lua/orgmode/org/mappings.lua
@@ -165,7 +165,7 @@ function OrgMappings:toggle_checkbox()
   vim.fn.setline('.', new_line)
   local headline = Headline:new(tree_utils.closest_headline())
   local list = List:new(tree_utils.closest_list())
-  headline:update_checkbox_status(list)
+  headline:update_cookie(list)
 end
 
 function OrgMappings:timestamp_up_day()
diff --git a/lua/orgmode/treesitter/headline.lua b/lua/orgmode/treesitter/headline.lua
index b0cd9f7a9..c02014547 100644
--- a/lua/orgmode/treesitter/headline.lua
+++ b/lua/orgmode/treesitter/headline.lua
@@ -120,12 +120,12 @@ function Headline:remove_closed_date()
   tree_utils.set_node_text(dates['CLOSED'], '', true)
 end
 
-function Headline:checkbox_status()
+function Headline:cookie()
   return self:parse('%[%d?/%d?%]')
 end
 
-function Headline:update_checkbox_status(list)
-  local checkbox_status = self:checkbox_status()
+function Headline:update_cookie(list)
+  local checkbox_status = self:cookie()
   if not checkbox_status then
     return
   end
diff --git a/tests/plenary/ui/mappings_spec.lua b/tests/plenary/ui/mappings_spec.lua
index e6fb3086c..9cf763bf1 100644
--- a/tests/plenary/ui/mappings_spec.lua
+++ b/tests/plenary/ui/mappings_spec.lua
@@ -1597,7 +1597,7 @@ describe('Mappings', function()
     }, vim.api.nvim_buf_get_lines(0, 0, 5, false))
   end)
 
-  it('should update the checklist status on a headline', function()
+  it('should update the checklist cookies on a headline', function()
     helpers.load_file_content({
       '* Test orgmode [/]',
       '- [ ] checkbox item',

From 71f38a7a801f7cdb41d8fe71274ac87699bc9fa7 Mon Sep 17 00:00:00 2001
From: Tray Dennis <tman1300@aol.com>
Date: Mon, 9 May 2022 21:08:08 -0400
Subject: [PATCH 3/6] Handle nested lists with cookies

---
 lua/orgmode/org/mappings.lua        |  2 +-
 lua/orgmode/treesitter/headline.lua | 14 ------------
 lua/orgmode/treesitter/list.lua     | 35 +++++++++++++++++++++++++++++
 tests/plenary/ui/mappings_spec.lua  | 11 +++++++++
 4 files changed, 47 insertions(+), 15 deletions(-)

diff --git a/lua/orgmode/org/mappings.lua b/lua/orgmode/org/mappings.lua
index 529b3f98d..95d8bd976 100644
--- a/lua/orgmode/org/mappings.lua
+++ b/lua/orgmode/org/mappings.lua
@@ -165,7 +165,7 @@ function OrgMappings:toggle_checkbox()
   vim.fn.setline('.', new_line)
   local headline = Headline:new(tree_utils.closest_headline())
   local list = List:new(tree_utils.closest_list())
-  headline:update_cookie(list)
+  list:update_parent_cookie()
 end
 
 function OrgMappings:timestamp_up_day()
diff --git a/lua/orgmode/treesitter/headline.lua b/lua/orgmode/treesitter/headline.lua
index c02014547..d2554ff8c 100644
--- a/lua/orgmode/treesitter/headline.lua
+++ b/lua/orgmode/treesitter/headline.lua
@@ -124,20 +124,6 @@ function Headline:cookie()
   return self:parse('%[%d?/%d?%]')
 end
 
-function Headline:update_cookie(list)
-  local checkbox_status = self:cookie()
-  if not checkbox_status then
-    return
-  end
-
-  local checkboxes = list:checkboxes()
-  local checked_boxes = vim.tbl_filter(function(box)
-    return box:match('%[%w%]')
-  end, checkboxes)
-  local new_status = ('[%d/%d]'):format(#checked_boxes, #checkboxes)
-  tree_utils.set_node_text(checkbox_status, new_status)
-end
-
 -- @return tsnode, string
 function Headline:parse(pattern)
   local match = ''
diff --git a/lua/orgmode/treesitter/list.lua b/lua/orgmode/treesitter/list.lua
index 6fcc761ac..ac987c323 100644
--- a/lua/orgmode/treesitter/list.lua
+++ b/lua/orgmode/treesitter/list.lua
@@ -1,6 +1,7 @@
 local ts_utils = require('nvim-treesitter.ts_utils')
 local tree_utils = require('orgmode.utils.treesitter')
 local query = vim.treesitter.query
+local Headline = require('orgmode.treesitter.headline')
 
 local List = {}
 
@@ -11,6 +12,40 @@ function List:new(list_node)
   return data
 end
 
+-- Updates the cookie of the immediate parent
+-- This always checks for a parent list first
+-- then for a headline.
+function List:parent_cookie()
+  local parent_list = tree_utils.find_list(self.list:parent())
+  if parent_list then
+    -- We only care about the cookie if it's at the top
+    local top_item = parent_list:named_child(0)
+    local content = top_item:field('contents')[1]
+    -- The cookie should be the last thing on the line
+    local cookie_node = content:named_child(content:named_child_count() - 1)
+    if query.get_node_text(cookie_node, 0):match('%[%d?/%d?%]') then
+      return cookie_node
+    end
+  end
+
+  local parent_header = Headline:new(tree_utils.closest_headline())
+  return parent_header:cookie()
+end
+
+function List:update_parent_cookie()
+  local parent_cookie = self:parent_cookie()
+  if not parent_cookie then
+    return
+  end
+
+  local checkboxes = self:checkboxes()
+  local checked_boxes = vim.tbl_filter(function(box)
+    return box:match('%[%w%]')
+  end, checkboxes)
+  local new_status = ('[%d/%d]'):format(#checked_boxes, #checkboxes)
+  tree_utils.set_node_text(parent_cookie, new_status)
+end
+
 function List:checkboxes()
   return vim.tbl_map(function(node)
     local text = query.get_node_text(node, 0)
diff --git a/tests/plenary/ui/mappings_spec.lua b/tests/plenary/ui/mappings_spec.lua
index 9cf763bf1..a87a8b86f 100644
--- a/tests/plenary/ui/mappings_spec.lua
+++ b/tests/plenary/ui/mappings_spec.lua
@@ -1607,4 +1607,15 @@ describe('Mappings', function()
     vim.cmd([[exe "norm \<C-space>"]])
     assert.are.same('* Test orgmode [1/2]', vim.fn.getline(1))
   end)
+
+  it('should update the checklist cookies on a parent list', function()
+    helpers.load_file_content({
+      '- Test orgmode [/]',
+      '  - [ ] checkbox item',
+      '  - [ ] checkbox item',
+    })
+    vim.fn.cursor(2, 1)
+    vim.cmd([[exe "norm \<C-space>"]])
+    assert.are.same('- Test orgmode [1/2]', vim.fn.getline(1))
+  end)
 end)

From 3194dad7f3e6b15312bc4d3d2846ad5ef52f51f2 Mon Sep 17 00:00:00 2001
From: Tray Dennis <tman1300@aol.com>
Date: Mon, 9 May 2022 21:35:45 -0400
Subject: [PATCH 4/6] Add support for percentage cookies

---
 lua/orgmode/treesitter/headline.lua |  6 +++++-
 lua/orgmode/treesitter/list.lua     | 11 +++++++++--
 lua/orgmode/utils/treesitter.lua    |  1 +
 tests/plenary/ui/mappings_spec.lua  | 23 +++++++++++++++++++++++
 4 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/lua/orgmode/treesitter/headline.lua b/lua/orgmode/treesitter/headline.lua
index d2554ff8c..2b7c3b066 100644
--- a/lua/orgmode/treesitter/headline.lua
+++ b/lua/orgmode/treesitter/headline.lua
@@ -121,7 +121,11 @@ function Headline:remove_closed_date()
 end
 
 function Headline:cookie()
-  return self:parse('%[%d?/%d?%]')
+  local cookie = self:parse('%[%d?/%d?%]')
+  if cookie then
+    return cookie
+  end
+  return self:parse('%[%d?%d?%d?%%%]')
 end
 
 -- @return tsnode, string
diff --git a/lua/orgmode/treesitter/list.lua b/lua/orgmode/treesitter/list.lua
index ac987c323..6f5d1b9a7 100644
--- a/lua/orgmode/treesitter/list.lua
+++ b/lua/orgmode/treesitter/list.lua
@@ -23,7 +23,9 @@ function List:parent_cookie()
     local content = top_item:field('contents')[1]
     -- The cookie should be the last thing on the line
     local cookie_node = content:named_child(content:named_child_count() - 1)
-    if query.get_node_text(cookie_node, 0):match('%[%d?/%d?%]') then
+
+    local text = query.get_node_text(cookie_node, 0)
+    if text:match('%[%d?/%d?%]') or text:match('%[%d?%d?%d?%%%]') then
       return cookie_node
     end
   end
@@ -42,7 +44,12 @@ function List:update_parent_cookie()
   local checked_boxes = vim.tbl_filter(function(box)
     return box:match('%[%w%]')
   end, checkboxes)
-  local new_status = ('[%d/%d]'):format(#checked_boxes, #checkboxes)
+  local new_status
+  if query.get_node_text(parent_cookie, 0):find('%%') then
+    new_status = ('[%d%%]'):format((#checked_boxes/#checkboxes) * 100)
+  else
+    new_status = ('[%d/%d]'):format(#checked_boxes, #checkboxes)
+  end
   tree_utils.set_node_text(parent_cookie, new_status)
 end
 
diff --git a/lua/orgmode/utils/treesitter.lua b/lua/orgmode/utils/treesitter.lua
index 8fdb08cff..4be2910f8 100644
--- a/lua/orgmode/utils/treesitter.lua
+++ b/lua/orgmode/utils/treesitter.lua
@@ -1,5 +1,6 @@
 local ts_utils = require('nvim-treesitter.ts_utils')
 local config = require('orgmode.config')
+local query = vim.treesitter.query
 local M = {}
 
 function M.current_node()
diff --git a/tests/plenary/ui/mappings_spec.lua b/tests/plenary/ui/mappings_spec.lua
index a87a8b86f..55c13b629 100644
--- a/tests/plenary/ui/mappings_spec.lua
+++ b/tests/plenary/ui/mappings_spec.lua
@@ -1618,4 +1618,27 @@ describe('Mappings', function()
     vim.cmd([[exe "norm \<C-space>"]])
     assert.are.same('- Test orgmode [1/2]', vim.fn.getline(1))
   end)
+
+  it('should update the checklist cookies with a percentage within a headline', function()
+    helpers.load_file_content({
+      '* Test orgmode [%]',
+      '- [ ] checkbox item',
+      '- [ ] checkbox item',
+    })
+    vim.fn.cursor(2, 1)
+    vim.cmd([[exe "norm \<C-space>"]])
+    assert.are.same('* Test orgmode [50%]', vim.fn.getline(1))
+  end)
+
+  it('should update the checklist cookies with a percentage within a nested list', function()
+    helpers.load_file_content({
+      '- Test orgmode [%]',
+      '  - [ ] checkbox item',
+      '  - [ ] checkbox item',
+      '  - [ ] checkbox item',
+    })
+    vim.fn.cursor(2, 1)
+    vim.cmd([[exe "norm \<C-space>"]])
+    assert.are.same('- Test orgmode [33%]', vim.fn.getline(1))
+  end)
 end)

From 003b9d91f97fd6941cf00173f2584c74d79f0679 Mon Sep 17 00:00:00 2001
From: Tray Dennis <tman1300@aol.com>
Date: Mon, 9 May 2022 21:36:05 -0400
Subject: [PATCH 5/6] Update formatting

---
 lua/orgmode/treesitter/list.lua | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lua/orgmode/treesitter/list.lua b/lua/orgmode/treesitter/list.lua
index 6f5d1b9a7..520233d55 100644
--- a/lua/orgmode/treesitter/list.lua
+++ b/lua/orgmode/treesitter/list.lua
@@ -46,7 +46,7 @@ function List:update_parent_cookie()
   end, checkboxes)
   local new_status
   if query.get_node_text(parent_cookie, 0):find('%%') then
-    new_status = ('[%d%%]'):format((#checked_boxes/#checkboxes) * 100)
+    new_status = ('[%d%%]'):format((#checked_boxes / #checkboxes) * 100)
   else
     new_status = ('[%d/%d]'):format(#checked_boxes, #checkboxes)
   end

From 34bc99733cdcb16ba4865dd6223493a3c85deff2 Mon Sep 17 00:00:00 2001
From: Tray Dennis <tman1300@aol.com>
Date: Tue, 10 May 2022 06:16:04 -0400
Subject: [PATCH 6/6] remove unused headline variable

---
 lua/orgmode/org/mappings.lua | 1 -
 1 file changed, 1 deletion(-)

diff --git a/lua/orgmode/org/mappings.lua b/lua/orgmode/org/mappings.lua
index 95d8bd976..f7db77e0b 100644
--- a/lua/orgmode/org/mappings.lua
+++ b/lua/orgmode/org/mappings.lua
@@ -163,7 +163,6 @@ function OrgMappings:toggle_checkbox()
   checkbox = checkbox:gsub('%[[%sXx%-]?%]$', new_val)
   local new_line = line:gsub(pattern, checkbox)
   vim.fn.setline('.', new_line)
-  local headline = Headline:new(tree_utils.closest_headline())
   local list = List:new(tree_utils.closest_list())
   list:update_parent_cookie()
 end