feature: add support for cookies

fix nvim-orgmode#81

lua/orgmode/init.lua

@@ -55,6 +55,9 @@ function Org:setup_autocmds()
   )
   vim.cmd([[autocmd FileType org call luaeval('require("orgmode").reload(_A)', expand('<afile>:p'))]])
   vim.cmd([[autocmd CursorHold,CursorHoldI *.org,*.org_archive lua require('orgmode.org.diagnostics').report()]])
+  vim.cmd(
+    [[autocmd BufWinEnter,FileChangedShellPost,TextChanged,TextChangedI *.org,*.org_archive lua require('orgmode.org.syntax').add_cookies()]]
+  )
   vim.cmd([[augroup END]])
   vim.cmd([[command! OrgDiagnostics lua require('orgmode.org.diagnostics').print()]])
 end

lua/orgmode/org/syntax.lua

@@ -1,5 +1,11 @@
 local Files = require('orgmode.parser.files')
 local config = require('orgmode.config')
+local utils = require('orgmode.utils')
+
+local cookie_namespace = vim.api.nvim_create_namespace('orgmode-cookies')
+local cookie_query = vim.treesitter.parse_query('org', '(cookie) @cookie')
+local checkbox_query = vim.treesitter.parse_query('org', '(list (listitem (checkbox) @checkbox))')
+local list_in_body_query = vim.treesitter.parse_query('org', '(body (list) @list)')
 
 local function load_code_blocks()
   local file = vim.api.nvim_buf_get_name(0)
@@ -36,7 +42,73 @@ local function add_todo_keywords_to_spellgood()
   end
 end
 
+local function _display_cookie_virt_text(cookie, checked, total)
+  local row, column = cookie:start()
+  local cookie_text = vim.treesitter.get_node_text(cookie, 0)
+
+  -- TODO: fix this in treesitter
+  local _, leading_spaces = cookie_text:find('^%s+')
+  cookie_text = vim.trim(cookie_text)
+
+  local text
+  if cookie_text == '[/]' then
+    text = string.format('[%d/%d]', checked, total)
+  else
+    text = string.format('[%d%%]', (100 * checked) / total)
+  end
+  vim.api.nvim_buf_set_extmark(0, cookie_namespace, row, column + (leading_spaces or 0), {
+    -- TODO: make highlight group configurable
+    virt_text = { { text, 'Error' } },
+    virt_text_pos = 'overlay',
+    hl_mode = 'combine',
+  })
+end
+
+local function _get_cookie_checked_and_total(parent)
+  local parent_type = parent:type()
+  local start_row, _, end_row, _ = parent:range()
+  local checked, total = 0, 0
+  for _, checkbox in checkbox_query:iter_captures(parent, 0, start_row, end_row + 1) do
+    local closest_parent = utils.get_closest_parent_of_type(checkbox, parent_type)
+    if closest_parent and closest_parent == parent then -- only count direct children
+      local checkbox_text = vim.treesitter.get_node_text(checkbox, 0)
+      if checkbox_text:match('%[[x|X]%]') then
+        checked = checked + 1
+      end
+      total = total + 1
+    end
+  end
+
+  return checked, total
+end
+
+local function add_cookies()
+  local parser = vim.treesitter.get_parser(0, 'org')
+  local root = parser:parse()[1]:root()
+
+  vim.api.nvim_buf_clear_namespace(0, cookie_namespace, 0, -1)
+
+  for _, cookie in cookie_query:iter_captures(root, 0, 0, -1) do
+    local parent = cookie:parent()
+    local checked, total = 0, 0
+    if parent:type() == 'headline' then
+      parent = parent:parent()
+      local start_row, _, end_row, _ = parent:range()
+      for _, list in list_in_body_query:iter_captures(parent, 0, start_row, end_row + 1) do
+        local c_checked, c_total = _get_cookie_checked_and_total(list)
+        checked = checked + c_checked
+        total = total + c_total
+      end
+    else
+      checked, total = _get_cookie_checked_and_total(parent)
+    end
+
+    _display_cookie_virt_text(cookie, checked, total)
+  end
+end
+
 return {
   load_code_blocks = load_code_blocks,
   add_todo_keywords_to_spellgood = add_todo_keywords_to_spellgood,
+  add_cookies = add_cookies,
 }