feature: add support for cookies

lua/orgmode/org/checkboxes.lua

@@ -0,0 +1,132 @@
+local utils = require('orgmode.utils')
+local ts_utils = require('nvim-treesitter.ts_utils')
+local headline_cookie_query = vim.treesitter.parse_query('org', '(headline (cookie) @cookie)')
+
+local checkboxes = {}
+
+---@param parent userdata
+local function _get_checked_and_total_checkboxes(parent)
+  local checked, total = 0, 0
+  for child in parent:iter_children() do
+    if child:type() == 'listitem' then
+      for listitem_child in child:iter_children() do
+        if listitem_child:type() == 'checkbox' then
+          local checkbox_text = vim.treesitter.get_node_text(listitem_child, 0)
+          if checkbox_text:match('%[[x|X]%]') then
+            checked = checked + 1
+          end
+          total = total + 1
+        end
+      end
+    end
+  end
+  return checked, total
+end
+
+---@param action string [on|off|toggle|children]
+---@param checkbox userdata
+---@param checked_children number
+---@param total_children number
+local function _update_checkbox_text(action, checkbox, checked_children, total_children)
+  local checkbox_text
+  if action == 'on' then
+    checkbox_text = '[X]'
+  elseif action == 'off' then
+    checkbox_text = '[ ]'
+  elseif action == 'toggle' then
+    checkbox_text = vim.treesitter.get_node_text(checkbox, 0)
+    if checkbox_text:match('%[[xX]%]') then
+      checkbox_text = '[ ]'
+    else
+      checkbox_text = '[X]'
+    end
+  elseif action == 'children' then
+    checkbox_text = '[ ]'
+    if checked_children == total_children then
+      checkbox_text = '[x]'
+    elseif checked_children > 0 then
+      checkbox_text = '[-]'
+    end
+  end
+
+  if checkbox_text then
+    utils.update_node_text(checkbox, { checkbox_text })
+  end
+end
+
+---@param cookie userdata
+---@param checked_children number
+---@param total_children number
+local function _update_cookie_text(cookie, checked_children, total_children)
+  local cookie_text = vim.treesitter.get_node_text(cookie, 0)
+
+  if total_children == nil then
+    checked_children, total_children = 0, 0
+  end
+
+  local new_cookie
+  if cookie_text:find('/') then
+    new_cookie = string.format('[%d/%d]', checked_children, total_children)
+  else
+    if total_children > 0 then
+      new_cookie = string.format('[%d%%%%]', (100 * checked_children) / total_children)
+    else
+      new_cookie = '[0%%%]'
+    end
+  end
+  cookie_text = cookie_text:gsub('%[.*%]', new_cookie)
+  utils.update_node_text(cookie, { cookie_text })
+end
+
+---@param action string [on|off|toggle|children]
+---@param node userdata
+---@param checked_children number
+---@param total_children number
+function checkboxes.update_checkbox(action, node, checked_children, total_children)
+  if not node then
+    node = utils.get_closest_parent_of_type(ts_utils.get_node_at_cursor(0), 'listitem')
+    if not node then
+      return
+    end
+  end
+
+  local checkbox
+  local cookie
+  for child in node:iter_children() do
+    if child:type() == 'checkbox' then
+      checkbox = child
+    elseif child:type() == 'itemtext' then
+      local c_child = child:named_child(0)
+      if c_child and c_child:type() == 'cookie' then
+        cookie = c_child
+      end
+    end
+  end
+
+  if checkbox then
+    _update_checkbox_text(action, checkbox, checked_children, total_children)
+  end
+
+  if cookie then
+    _update_cookie_text(cookie, checked_children, total_children)
+  end
+
+  local listitem_parent = utils.get_closest_parent_of_type(node:parent(), 'listitem')
+  if listitem_parent then
+    local list_parent = utils.get_closest_parent_of_type(node, 'list')
+    local checked, total = _get_checked_and_total_checkboxes(list_parent)
+    return checkboxes.update_checkbox('children', listitem_parent, checked, total)
+  end
+
+  local section = utils.get_closest_parent_of_type(node:parent(), 'section')
+  if section then
+    local list_parent = utils.get_closest_parent_of_type(node, 'list')
+    local checked, total = _get_checked_and_total_checkboxes(list_parent)
+    local start_row, _, end_row, _ = section:range()
+    for _, headline_cookie in headline_cookie_query:iter_captures(section, 0, start_row, end_row + 1) do
+      _update_cookie_text(headline_cookie, checked, total)
+    end
+  end
+end
+
+return checkboxes

lua/orgmode/org/mappings.lua

@@ -8,6 +8,7 @@ local utils = require('orgmode.utils')
 local Files = require('orgmode.parser.files')
 local config = require('orgmode.config')
 local Help = require('orgmode.objects.help')
+local checkboxes = require('orgmode.org.checkboxes')
 
 ---@class OrgMappings
 ---@field capture Capture
@@ -142,18 +143,8 @@ function OrgMappings:global_cycle()
   return vim.cmd([[silent! norm!zx]])
 end
 
--- TODO: Add hierarchy
 function OrgMappings:toggle_checkbox()
-  local line = vim.fn.getline('.')
-  local pattern = '^(%s*[%-%+]%s*%[([%sXx%-]?)%])'
-  local checkbox, state = line:match(pattern)
-  if not checkbox then
-    return
-  end
-  local new_val = vim.trim(state) == '' and '[X]' or '[ ]'
-  checkbox = checkbox:gsub('%[[%sXx%-]?%]$', new_val)
-  local new_line = line:gsub(pattern, checkbox)
-  vim.fn.setline('.', new_line)
+  checkboxes.update_checkbox('toggle')
 end
 
 function OrgMappings:timestamp_up_day()
@@ -378,12 +369,12 @@ function OrgMappings:handle_return(suffix)
     return vim.cmd([[startinsert!]])
   end
 
-  if item.type == 'list' or item.type == 'listitem' then
+  if vim.tbl_contains({ 'list', 'listitem', 'cookie' }, item.type) then
     vim.cmd([[normal! ^]])
     item = Files.get_current_file():get_current_node()
   end
 
-  if item.type == 'itemtext' or item.type == 'bullet' or item.type == 'checkbox' or item.type == 'description' then
+  if vim.tbl_contains({ 'itemtext', 'bullet', 'checkbox', 'description' }, item.type) then
     local list_item = item.node:parent()
     if list_item:type() ~= 'listitem' then
       return
@@ -442,6 +433,7 @@ function OrgMappings:handle_return(suffix)
       vim.lsp.util.apply_text_edits(text_edits, 0)
 
       vim.fn.cursor(end_row + 2 + (add_empty_line and 1 or 0), 0) -- +1 for 0 index and +1 for next line
+      checkboxes.update_checkbox('off', ts_utils.get_next_node(list_item))
       vim.cmd([[startinsert!]])
     end
   end

lua/orgmode/utils/init.lua

@@ -321,6 +321,13 @@ function utils.get_node_text(node, content)
   end
 end
 
+---@param node userdata
+---@param text string[]
+function utils.update_node_text(node, text)
+  local start_row, start_col, end_row, end_col = node:range()
+  vim.api.nvim_buf_set_text(0, start_row, start_col, end_row, end_col, text)
+end
+
 ---@param node table
 ---@param type string
 ---@return table