Add support for nested cookies

Author: lukas-reineke

lua/orgmode/org/mappings.lua

@@ -12,7 +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')
+local Listitem = require('orgmode.treesitter.listitem')
 
 ---@class OrgMappings
 ---@field capture Capture
@@ -38,8 +38,6 @@ function OrgMappings:archive()
     return utils.echo_warning('This file is already an archive file.')
   end
   local item = file:get_closest_headline()
-  file = Files.get_current_file()
-  item = file:get_closest_headline()
   local archive_location = file:get_archive_file_location()
   local archive_directory = vim.fn.fnamemodify(archive_location, ':p:h')
   if vim.fn.isdirectory(archive_directory) == 0 then
@@ -156,20 +154,18 @@ 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
+  local win_view = vim.fn.winsaveview()
+  -- move to the first non-blank character so the current treesitter node is the listitem
+  vim.cmd([[normal! _]])
+
+  vim.treesitter.get_parser(0, 'org'):parse()
+  local listitem = tree_utils.find_parent_type(tree_utils.current_node(), 'listitem')
+  if listitem then
+    Listitem:new(listitem):update_checkbox('toggle')
   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)
-  local list = List:new(tree_utils.closest_list())
-  list:update_parent_cookie()
+
+  vim.fn.winrestview(win_view)
 end
 
 function OrgMappings:timestamp_up_day()
@@ -463,12 +459,12 @@ function OrgMappings:handle_return(suffix)
   end
 
   if item.type == 'paragraph' or item.type == 'bullet' then
-    local list_item = item.node:parent()
-    if list_item:type() ~= 'listitem' then
+    local listitem = item.node:parent()
+    if listitem:type() ~= 'listitem' then
       return
     end
-    local line = vim.fn.getline(list_item:start() + 1)
-    local end_row, _ = list_item:end_()
+    local line = vim.fn.getline(listitem:start() + 1)
+    local end_row, _ = listitem:end_()
     local next_line_node = current_file:get_node_at_cursor({ end_row + 1, 0 })
     local second_line_node = current_file:get_node_at_cursor({ end_row + 2, 0 })
     local is_end_of_file = next_line_node
@@ -503,7 +499,7 @@ function OrgMappings:handle_return(suffix)
         newText = plain_list .. ' \n',
       })
     elseif number_in_list then
-      local next_sibling = list_item
+      local next_sibling = listitem
       local counter = 1
       while next_sibling do
         local bullet = next_sibling:child(0)
@@ -531,6 +527,13 @@ function OrgMappings:handle_return(suffix)
       vim.lsp.util.apply_text_edits(text_edits, 0, constants.default_offset_encoding)
 
       vim.fn.cursor(end_row + 1 + (add_empty_line and 1 or 0), 0) -- +1 for next line
+
+      -- update all parents when we insert a new checkbox
+      if checkbox then
+        local new_listitem = tree_utils.find_parent_type(tree_utils.current_node(), 'listitem')
+        Listitem:new(new_listitem):update_checkbox('off')
+      end
+
       vim.cmd([[startinsert!]])
     end
   end

lua/orgmode/treesitter/headline.lua

@@ -191,13 +191,38 @@ function Headline:remove_closed_date()
 end
 
 function Headline:cookie()
-  local cookie = self:parse('%[%d?/%d?%]')
+  local cookie = self:parse('%[%d*/%d*%]')
   if cookie then
     return cookie
   end
   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 query.get_node_text(cookie, 0):find('%%') then
+      new_cookie_val = ('[%d%%]'):format((#checked_boxes / #total_boxes) * 100)
+    else
+      new_cookie_val = ('[%d/%d]'):format(#checked_boxes, #total_boxes)
+    end
+    tree_utils.set_node_text(cookie, new_cookie_val)
+  end
+end
+
+function Headline:child_checkboxes(list_node)
+  return vim.tbl_map(function(node)
+    local text = query.get_node_text(node, 0)
+    return text:match('%[.%]')
+  end, ts_utils.get_named_children(list_node))
+end
+
 -- @return tsnode, string
 function Headline:parse(pattern)
   local match = ''

lua/orgmode/treesitter/list.lua

@@ -0,0 +1,123 @@
+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 Listitem = {}
+
+function Listitem:new(listitem_node)
+  local data = { listitem = listitem_node }
+  setmetatable(data, self)
+  self.__index = self
+  return data
+end
+
+function Listitem:get_new_checkbox_value(action, current_value, total_child_checkboxes, checked_child_checkboxes)
+  if action == 'on' then
+    return '[X]'
+  elseif action == 'off' then
+    return '[ ]'
+  elseif action == 'toggle' then
+    return (current_value == '[X]' or current_value == '[x]') and '[ ]' or '[X]'
+  elseif action == 'children' then
+    if #checked_child_checkboxes == 0 then
+      return '[ ]'
+    elseif #checked_child_checkboxes == #total_child_checkboxes then
+      return '[X]'
+    end
+  end
+  return '[-]'
+end
+
+function Listitem:checkbox()
+  local contents = self.listitem:field('contents')
+  for _, content in ipairs(contents) do
+    if content:type() == 'paragraph' then
+      for child in content:iter_children() do
+        local text = query.get_node_text(child, 0)
+        if text:match('%[[Xx-]%]') then
+          return { text = text, range = { child:range() } }
+        end
+
+        -- empty checkboxes are split into two nodes, so we need to hack this together
+        local next_sibling = child:next_sibling()
+        if next_sibling and text == '[' and query.get_node_text(next_sibling, 0) == ']' then
+          local sr, sc, er, ec = child:range()
+          return { text = '[ ]', range = { sr, sc, er, ec + 2 } }
+        end
+      end
+    end
+  end
+end
+
+function Listitem:update_checkbox(action)
+  action = action or 'toggle'
+
+  local checkbox = self:checkbox()
+  local total_child_checkboxes = self:child_checkboxes() or {}
+  local checked_child_checkboxes = vim.tbl_filter(function(box)
+    return box:match('%[%w%]')
+  end, total_child_checkboxes)
+
+  if checkbox then
+    vim.api.nvim_buf_set_text(
+      0,
+      checkbox.range[1],
+      checkbox.range[2],
+      checkbox.range[3],
+      checkbox.range[4],
+      { self:get_new_checkbox_value(action, checkbox.text, total_child_checkboxes, checked_child_checkboxes) }
+    )
+  end
+
+  self:update_cookie(total_child_checkboxes, checked_child_checkboxes)
+
+  local parent_list = tree_utils.find_parent_type(self.listitem, 'list')
+  local parent_listitem = tree_utils.find_parent_type(parent_list, 'listitem')
+  if parent_listitem then
+    Listitem:new(parent_listitem):update_checkbox('children')
+  else
+    local parent_headline = tree_utils.closest_headline()
+    if parent_headline then
+      Headline:new(parent_headline):update_cookie(parent_list)
+    end
+  end
+end
+
+function Listitem:child_checkboxes()
+  local contents = self.listitem:field('contents')
+  for _, content in ipairs(contents) do
+    if content:type() == 'list' then
+      return vim.tbl_map(function(node)
+        local text = query.get_node_text(node, 0)
+        return text:match('%[.%]')
+      end, ts_utils.get_named_children(content))
+    end
+  end
+end
+
+function Listitem:cookie()
+  local content = self.listitem:field('contents')[1]
+  -- The cookie should be the last thing on the line
+  local cookie_node = content:named_child(content:named_child_count() - 1)
+
+  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
+
+function Listitem:update_cookie(total_child_checkboxes, checked_child_checkboxes)
+  local cookie = self:cookie()
+  if cookie then
+    local new_cookie_val
+    if query.get_node_text(cookie, 0):find('%%') then
+      new_cookie_val = ('[%d%%]'):format((#checked_child_checkboxes / #total_child_checkboxes) * 100)
+    else
+      new_cookie_val = ('[%d/%d]'):format(#checked_child_checkboxes, #total_child_checkboxes)
+    end
+    tree_utils.set_node_text(cookie, new_cookie_val)
+  end
+end
+
+return Listitem

lua/orgmode/treesitter/listitem.lua

@@ -1,6 +1,4 @@
 local ts_utils = require('nvim-treesitter.ts_utils')
-local config = require('orgmode.config')
-local query = vim.treesitter.query
 local M = {}
 
 function M.current_node()
@@ -28,23 +26,19 @@ function M.closest_headline()
   return M.find_headline(M.current_node())
 end
 
-function M.find_list(node)
-  if node:type() == 'list' then
+function M.find_parent_type(node, type)
+  if node:type() == type 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())
+  return M.find_parent_type(node:parent(), type)
 end
 
 -- @param front_trim boolean
 function M.set_node_text(node, text, front_trim)
+  local lines = vim.split(text, '\n', true)
   local sr, sc, er, ec = node:range()
   if string.len(text) == 0 then
     if front_trim then
@@ -53,7 +47,7 @@ function M.set_node_text(node, text, front_trim)
       ec = ec + 1
     end
   end
-  vim.api.nvim_buf_set_text(0, sr, sc, er, ec, { text })
+  pcall(vim.api.nvim_buf_set_text, 0, sr, sc, er, ec, lines)
 end
 
 return M