Merge pull request #272 from TravonteD/tree-sitter-actions

Use treesitter for buffer actions

Author: kristijanhusak

lua/orgmode/org/mappings.lua

@@ -10,6 +10,8 @@ local config = require('orgmode.config')
 local constants = require('orgmode.utils.constants')
 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')
 
 ---@class OrgMappings
 ---@field capture Capture
@@ -307,21 +309,20 @@ function OrgMappings:priority_down()
 end
 
 function OrgMappings:set_priority(direction)
-  local item = Files.get_closest_headline()
-  local priority = PriorityState:new(item.priority)
+  local headline = Headline:new(tree_utils.closest_headline())
+  local _, current_priority = headline:priority()
+  local priority_state = PriorityState:new(current_priority)
 
   local new_priority
-  if direction then
-    new_priority = direction == 'up' and priority:increase() or priority:decrease()
+  if direction == 'up' then
+    new_priority = priority_state:increase()
+  elseif direction == 'down' then
+    new_priority = priority_state:decrease()
   else
-    new_priority = priority:prompt_user()
-  end
-
-  if not new_priority then
-    return
+    new_priority = priority_state:prompt_user()
   end
 
-  item:set_priority(new_priority)
+  headline:set_priority(new_priority)
 end
 
 function OrgMappings:todo_next_state()
@@ -363,9 +364,8 @@ function OrgMappings:toggle_heading()
 end
 
 function OrgMappings:_todo_change_state(direction)
-  local item = Files.get_closest_headline()
-  local was_done = item:is_done()
-  local old_state = item.todo_keyword.value
+  local headline = Headline:new(tree_utils.closest_headline())
+  local _, old_state, was_done = headline:todo()
   local changed = self:_change_todo_state(direction, true)
   if not changed then
     return
@@ -379,10 +379,10 @@ function OrgMappings:_todo_change_state(direction)
   if #repeater_dates == 0 then
     local log_time = config.org_log_done == 'time'
     if log_time and item:is_done() and not was_done then
-      item:add_closed_date()
+      headline:add_closed_date()
     end
     if log_time and not item:is_done() and was_done then
-      item:remove_closed_date()
+      headline:remove_closed_date()
     end
     return item
   end
@@ -753,9 +753,9 @@ end
 ---@param use_fast_access boolean
 ---@return string
 function OrgMappings:_change_todo_state(direction, use_fast_access)
-  local item = Files.get_closest_headline()
-  local todo = item.todo_keyword
-  local todo_state = TodoState:new({ current_state = todo.value })
+  local headline = Headline:new(tree_utils.closest_headline())
+  local todo, current_keyword = headline:todo()
+  local todo_state = TodoState:new({ current_state = current_keyword })
   local next_state = nil
   if use_fast_access and todo_state:has_fast_access() then
     next_state = todo_state:open_fast_access()
@@ -773,25 +773,14 @@ function OrgMappings:_change_todo_state(direction, use_fast_access)
     return false
   end
 
-  if next_state.value == todo.value then
+  if next_state.value == current_keyword then
     if todo.value ~= '' then
       utils.echo_info('TODO state was already ', { { next_state.value, next_state.hl } })
     end
     return false
   end
 
-  local linenr = item.range.start_line
-  local stars = string.rep('%*', item.level)
-  local old_state = vim.pesc(todo.value)
-  if old_state ~= '' then
-    old_state = old_state .. '%s+'
-  end
-  local new_state = next_state.value
-  if new_state ~= '' then
-    new_state = new_state .. ' '
-  end
-  local new_line = vim.fn.getline(linenr):gsub('^' .. stars .. '%s+' .. old_state, stars .. ' ' .. new_state)
-  vim.fn.setline(linenr, new_line)
+  headline:set_todo(next_state.value)
   return true
 end
 

lua/orgmode/treesitter/headline.lua

@@ -0,0 +1,137 @@
+local ts_utils = require('nvim-treesitter.ts_utils')
+local tree_utils = require('orgmode.utils.treesitter')
+local config = require('orgmode.config')
+local query = vim.treesitter.query
+
+local Headline = {}
+
+---@param headline userdata tree sitter headline node
+function Headline:new(headline_node)
+  local data = { headline = headline_node }
+  setmetatable(data, self)
+  self.__index = self
+  return data
+end
+
+function Headline:stars()
+  return self.headline:field('stars')[1]
+end
+
+function Headline:priority()
+  return self:parse('%[#(%w+)%]')
+end
+
+function Headline:set_priority(priority)
+  local current_priority = self:priority()
+  if current_priority then
+    local text = (vim.trim(priority) == '') and '' or ('[#%s]'):format(priority)
+    tree_utils.set_node_text(current_priority, text)
+    return
+  end
+
+  local todo = self:todo()
+  if todo then
+    local text = query.get_node_text(todo, 0)
+    tree_utils.set_node_text(todo, ('%s [#%s]'):format(text, priority))
+    return
+  end
+
+  local stars = self:stars()
+  local text = query.get_node_text(stars, 0)
+  tree_utils.set_node_text(stars, ('%s [#%s]'):format(text, priority))
+end
+
+function Headline:set_todo(keyword)
+  local current_todo = self:todo()
+  if current_todo then
+    tree_utils.set_node_text(current_todo, keyword)
+    return
+  end
+
+  local stars = self:stars()
+  local text = query.get_node_text(stars, 0)
+  tree_utils.set_node_text(stars, string.format('%s %s', text, keyword))
+end
+
+function Headline:item()
+  return self.headline:field('item')[1]
+end
+
+-- Returns the headlines todo node, it's keyword,
+-- and if it's in done state
+-- @return Node, string, boolean
+function Headline:todo()
+  local keywords = config.todo_keywords.ALL
+  local done_keywords = config.todo_keywords.DONE
+  for _, word in ipairs(keywords) do
+    local todo = self:parse(word:gsub('-', '%%-'))
+    if todo then
+      return todo, word, vim.tbl_contains(done_keywords, word)
+    end
+  end
+end
+
+function Headline:plan()
+  local section = self.headline:parent()
+  for _, node in ipairs(ts_utils.get_named_children(section)) do
+    if node:type() == 'plan' then
+      return node
+    end
+  end
+end
+
+function Headline:dates()
+  local plan = self:plan()
+  local dates = {}
+  for _, node in ipairs(ts_utils.get_named_children(plan)) do
+    local name = vim.treesitter.query.get_node_text(node:named_child(0), 0)
+    dates[name] = node
+  end
+  return dates
+end
+
+function Headline:repeater_dates()
+  return vim.tbl_filter(function(entry)
+    local timestamp = entry:field('timestamp')[1]
+    for _, node in ipairs(ts_utils.get_named_children(timestamp)) do
+      if node:type() == 'repeat' then
+        return true
+      end
+    end
+  end, self:dates())
+end
+
+function Headline:add_closed_date()
+  local dates = self:dates()
+  if vim.tbl_count(dates) == 0 or dates['CLOSED'] then
+    return
+  end
+  local last_child = dates['DEADLINE'] or dates['SCHEDULED']
+  local ptext = query.get_node_text(last_child, 0)
+  local text = ptext .. ' CLOSED: [' .. vim.fn.strftime('%Y-%m-%d %a %H:%M') .. ']'
+  tree_utils.set_node_text(last_child, text)
+end
+
+function Headline:remove_closed_date()
+  local dates = self:dates()
+  if vim.tbl_count(dates) == 0 or not dates['CLOSED'] then
+    return
+  end
+  tree_utils.set_node_text(dates['CLOSED'], '', true)
+end
+
+-- @return tsnode, string
+function Headline:parse(pattern)
+  local match = ''
+  local matching_nodes = vim.tbl_filter(function(node)
+    local text = query.get_node_text(node, 0) or ''
+    local m = text:match(pattern)
+    if m then
+      match = text:match(pattern)
+      return true
+    end
+  end, ts_utils.get_named_children(self:item()))
+  return matching_nodes[1], match
+end
+
+return Headline

lua/orgmode/utils/treesitter.lua

@@ -0,0 +1,38 @@
+local ts_utils = require('nvim-treesitter.ts_utils')
+local config = require('orgmode.config')
+local M = {}
+
+-- walks the tree to find a headline
+function M.find_headline(node)
+  if node:type() == 'headline' then
+    return node
+  elseif node:type() == 'section' then
+    -- The headline is always the first child of a section
+    return ts_utils.get_named_children(node)[1]
+  elseif node:parent() then
+    return M.find_headline(node:parent())
+  else
+    return nil
+  end
+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()))
+end
+
+-- @param front_trim boolean
+function M.set_node_text(node, text, front_trim)
+  local sr, sc, er, ec = node:range()
+  if string.len(text) == 0 then
+    if front_trim then
+      sc = sc - 1
+    else
+      ec = ec + 1
+    end
+  end
+  vim.api.nvim_buf_set_text(0, sr, sc, er, ec, { text })
+end
+
+return M