feat(treesitter): incremental selection

Co-authored-by: György Andorka <gyorgy.andorka@protonmail.com>

1 files changed, 583 insertions, 0 deletions

diff --git a/runtime/lua/vim/treesitter/_select.lua b/runtime/lua/vim/treesitter/_select.lua
new file mode 100644
index 0000000000..3800524c94
--- /dev/null
+++ b/runtime/lua/vim/treesitter/_select.lua
@@ -0,0 +1,583 @@
+local Range = require('vim.treesitter._range')
+
+--- This is (currently only) used for saving what child one is in when doing
+--- `select_parent` so that if they later `select_child` on the parent-node,
+--- they get back to the child-node they were in instead of the parents first
+--- child-node.
+---
+--- @type {[integer]:vim.treesitter.select.node,[any]:any}
+local history = {
+  --- @type integer?
+  bufnr = nil,
+
+  --- @type integer?
+  changedtick = nil,
+
+  --- @type string?
+  current_node_id = nil,
+}
+
+--- The reason for a wrapper around `TSNode` is because we need to store the
+--- information about which tstree-range they are in (as a tstree may be
+--- disjointed), where region is the return value of
+--- `TSTree:included_ranges(false)` with next to eachother ranges combined
+--- (e.g. {{0,0,1,1},{1,1,2,2}} -> {{0,0,2,2}}).
+---
+--- @class vim.treesitter.select.node
+--- @field node TSNode
+--- @field top vim.treesitter.select.node.top
+
+--- @class vim.treesitter.select.node.top: vim.treesitter.select.node
+--- @field ltree vim.treesitter.LanguageTree
+--- @field region Range4
+
+local M = {}
+
+--- @param node vim.treesitter.select.node
+--- @return string
+local function node_id(node)
+  return ('%s:%s'):format(table.concat({ unpack(node.top.region) }, ':'), node.node:id())
+end
+
+--- @param r1 Range4
+--- @param r2 Range4
+--- @return Range4?
+local function range_intersection(r1, r2)
+  if not Range.intercepts(r1, r2) then
+    return
+  end
+
+  local rs = Range.cmp_pos.le(r1[1], r1[2], r2[1], r2[2]) and r2 or r1
+  local re = Range.cmp_pos.ge(r1[3], r1[4], r2[3], r2[4]) and r2 or r1
+  return { rs[1], rs[2], re[3], re[4] }
+end
+
+--- @param r1 Range4
+--- @param r2 Range4
+--- @boolean
+local function range_is_same(r1, r2)
+  local srow_1, scol_1, erow_1, ecol_1 = Range.unpack4(r1)
+  local srow_2, scol_2, erow_2, ecol_2 = Range.unpack4(r2)
+  return srow_1 == srow_2 and scol_1 == scol_2 and erow_1 == erow_2 and ecol_1 == ecol_2
+end
+
+--- @param node vim.treesitter.select.node
+--- @return Range4
+local function node_range(node)
+  local node_range_ = { node.node:range() }
+
+  return range_intersection(node.top.region, node_range_) or { 0, 0, 0, 0 }
+end
+
+--- @param node1 vim.treesitter.select.node
+--- @param node2 vim.treesitter.select.node
+--- @return boolean
+local function node_is_same_range(node1, node2)
+  return range_is_same(node_range(node1), node_range(node2))
+end
+
+--- @param node vim.treesitter.select.node
+--- @return boolean
+local function node_is_size_0(node)
+  local srow, scol, erow, ecol = Range.unpack4(node_range(node))
+  return srow == erow and scol == ecol
+end
+
+--- @param tsnode TSNode
+--- @param relative vim.treesitter.select.node
+--- @return vim.treesitter.select.node
+local function create_node(tsnode, relative)
+  assert(tsnode:tree():root():equal(relative.top.node))
+
+  --- @type vim.treesitter.select.node
+  return {
+    node = tsnode,
+    top = relative.top,
+  }
+end
+
+--- @param tree TSTree
+--- @return Range4[]
+local function tree_get_ranges(tree)
+  --- @type Range4[]
+  local regions = {}
+  for _, tree_range in ipairs(tree:included_ranges(false)) do
+    local prev_region = regions[#regions]
+
+    if prev_region and prev_region[3] == tree_range[1] and prev_region[4] == tree_range[2] then
+      regions[#regions] = { prev_region[1], prev_region[2], tree_range[3], tree_range[4] }
+    else
+      table.insert(regions, tree_range)
+    end
+  end
+  return regions
+end
+
+--- @param tree TSTree
+--- @param region Range4
+--- @param ltree vim.treesitter.LanguageTree
+--- @return vim.treesitter.select.node.top
+local function create_top_node(tree, region, ltree)
+  --- @type vim.treesitter.select.node.top
+  local self = {
+    node = tree:root(),
+    top = {} --[[@as any]],
+    ltree = ltree,
+    region = region,
+  }
+  self.top = self
+  return self
+end
+
+--- @param node1 vim.treesitter.select.node.top
+--- @param node2 vim.treesitter.select.node.top
+--- @return boolean
+local function top_node_is_higher_priority(node1, node2)
+  local srow1, scol1, erow1, ecol1 = Range.unpack4(node_range(node1))
+  local srow2, scol2, erow2, ecol2 = Range.unpack4(node_range(node2))
+
+  if M.TEST_SWITCH_PRIORITY then
+    if Range.cmp_pos.ne(srow1, scol1, srow2, scol2) then
+      return Range.cmp_pos.lt(srow1, scol1, srow2, scol2)
+    elseif Range.cmp_pos.ne(erow1, ecol1, erow2, ecol2) then
+      return Range.cmp_pos.lt(erow1, ecol1, erow2, ecol2)
+    elseif node1.ltree:lang() ~= node2.ltree:lang() then
+      return node1.ltree:lang() > node2.ltree:lang()
+    end
+    return node1.node:id() > node2.node:id()
+  else
+    if Range.cmp_pos.ne(srow1, scol1, srow2, scol2) then
+      return Range.cmp_pos.gt(srow1, scol1, srow2, scol2)
+    elseif Range.cmp_pos.ne(erow1, ecol1, erow2, ecol2) then
+      return Range.cmp_pos.gt(erow1, ecol1, erow2, ecol2)
+    elseif node1.ltree:lang() ~= node2.ltree:lang() then
+      return node1.ltree:lang() < node2.ltree:lang()
+    end
+    return node1.node:id() < node2.node:id()
+  end
+end
+
+--- @param range Range4
+--- @param top_node vim.treesitter.select.node.top?
+--- @param parent_chain vim.treesitter.select.node[]?
+--- @return vim.treesitter.select.node|false|nil nil: no parser, false: outside of root-node
+--- @return vim.treesitter.select.node[] either `parent_chain` or `alternative_nodes`
+local function get_node(range, top_node, parent_chain)
+  parent_chain = parent_chain or {}
+
+  if not top_node then
+    local parser = vim.treesitter.get_parser(nil, nil, { error = false })
+    if not parser then
+      return nil, {}
+    end
+
+    local tree = assert(parser:parse(range))[1]
+    top_node = create_top_node(tree, assert(tree:included_ranges(false)[1]), parser)
+
+    if not Range.contains(node_range(top_node), range) then
+      return false, { top_node } --[[alternative_nodes]]
+    end
+  end
+
+  assert(Range.contains(node_range(top_node), range))
+
+  --- @param node vim.treesitter.select.node|vim.treesitter.select.node.top
+  --- @return vim.treesitter.select.node|vim.treesitter.select.node.top
+  local function node_ignore_overlapped_handle_injection(node)
+    for _, child in pairs(top_node.ltree:children()) do
+      for _, child_tree in ipairs(child:trees()) do
+        for _, child_region in ipairs(tree_get_ranges(child_tree)) do
+          local child_root_node_range = { child_tree:root():range() }
+          local child_range = range_intersection(child_region, child_root_node_range)
+
+          local child_top_node = create_top_node(child_tree, child_region, child)
+          if
+            child_range
+            and Range.contains(child_range, range)
+            and (
+              not node.ltree
+              or top_node_is_higher_priority(
+                node --[[@as vim.treesitter.select.node.top]],
+                child_top_node
+              )
+            )
+          then
+            return node_ignore_overlapped_handle_injection(child_top_node)
+          elseif child_range and Range.intercepts(node_range(node), child_range) then
+            local child_parent_tsnode =
+              assert(top_node.node:named_descendant_for_range(unpack(child_range)))
+
+            if
+              (not node.ltree and vim.treesitter.is_ancestor(child_parent_tsnode, node.node))
+              or (
+                node.ltree
+                and top_node_is_higher_priority(
+                  node --[[@as vim.treesitter.select.node.top]],
+                  child_top_node
+                )
+              )
+            then
+              return create_node(child_parent_tsnode, top_node)
+            end
+          end
+        end
+      end
+    end
+
+    return node
+  end
+
+  local tsnode = assert(top_node.node:named_descendant_for_range(unpack(range)))
+  local node = create_node(tsnode, top_node)
+
+  node = node_ignore_overlapped_handle_injection(node)
+  if node.ltree then
+    local root_node_range = { node.node:range() }
+    local tree_range = node.top.region
+    local actual_range = assert(range_intersection(tree_range, root_node_range))
+    local parent_tsnode = assert(top_node.node:named_descendant_for_range(unpack(actual_range)))
+    table.insert(parent_chain, create_node(parent_tsnode, top_node))
+
+    --- @cast node vim.treesitter.select.node.top
+    return get_node(range, node, parent_chain), parent_chain
+  end
+  --- @cast node vim.treesitter.select.node
+
+  return node, parent_chain
+end
+
+--- @param node vim.treesitter.select.node
+--- @param parent_chain vim.treesitter.select.node[]
+--- @nodiscard
+--- @return vim.treesitter.select.node?
+--- @return vim.treesitter.select.node.top?
+local function node_get_parent_no_normalize(node, parent_chain)
+  local parent = node.node:parent()
+  if parent then
+    return create_node(parent, node)
+  end
+
+  return table.remove(parent_chain)
+end
+
+--- @param node vim.treesitter.select.node
+--- @return vim.treesitter.select.node
+local function node_normalize_up(node, parent_chain)
+  while true do
+    local parent = node_get_parent_no_normalize(node, parent_chain)
+    if parent and node_is_same_range(parent, node) then
+      node = parent
+    else
+      table.insert(parent_chain, parent)
+
+      return node
+    end
+  end
+  --- @diagnostic disable-next-line: missing-return
+end
+
+--- @param nodes vim.treesitter.select.node[]
+--- @param node vim.treesitter.select.node.top
+local function insert_remove_overlapped(nodes, node)
+  local n = 1
+  while nodes[n] do
+    if Range.intercepts(node_range(nodes[n]), node_range(node)) then
+      if
+        not nodes
+          [n] --[[@as any]]
+          .ltree
+        or top_node_is_higher_priority(nodes[n] --[[@as vim.treesitter.select.node.top]], node)
+      then
+        table.remove(nodes, n)
+      else
+        return
+      end
+    else
+      local nrow, ncol, _, _ = Range.unpack4(node_range(nodes[n]))
+      local _, _, erow, ecol = Range.unpack4(node_range(node))
+      if Range.cmp_pos.le(erow, ecol, nrow, ncol) then
+        table.insert(nodes, n, node)
+        return
+      end
+
+      n = n + 1
+    end
+  end
+
+  table.insert(nodes, node)
+end
+
+--- @param node vim.treesitter.select.node
+--- @return vim.treesitter.select.node[]
+local function node_get_children_no_normalize(node)
+  --- @param child_ TSNode
+  --- @return vim.treesitter.select.node
+  local children = vim.tbl_map(function(child_)
+    return create_node(child_, node)
+  end, node.node:named_children())
+
+  node.top.ltree:parse(node_range(node))
+
+  for _, child in pairs(node.top.ltree:children()) do
+    for _, child_tree in ipairs(child:trees()) do
+      for _, child_region in ipairs(tree_get_ranges(child_tree)) do
+        local child_root_node_range = { child_tree:root():range() }
+        local child_range = range_intersection(child_region, child_root_node_range)
+
+        if child_range and Range.contains(node_range(node), child_range) then
+          local child_parent_tsnode =
+            assert(node.top.node:named_descendant_for_range(unpack(child_range)))
+
+          if node.node:equal(child_parent_tsnode) then
+            local child_node = create_top_node(child_tree, child_region, child)
+
+            insert_remove_overlapped(children, child_node)
+          end
+        end
+      end
+    end
+  end
+
+  return children
+end
+
+--- @param range Range4
+--- @param node vim.treesitter.select.node
+--- @return vim.treesitter.select.node?
+local function get_node_contained_in_range(range, node)
+  for _, child in ipairs(node_get_children_no_normalize(node)) do
+    if Range.contains(range, node_range(child)) and not node_is_size_0(child) then
+      return child
+    elseif Range.intercepts(range, node_range(child)) and not node_is_size_0(child) then
+      local smallest_node = get_node_contained_in_range(range, child)
+
+      if smallest_node then
+        return smallest_node
+      end
+    end
+  end
+end
+
+--- @param node vim.treesitter.select.node
+--- @return vim.treesitter.select.node
+local function node_normalize_down(node)
+  for _, child in ipairs(node_get_children_no_normalize(node)) do
+    if node_is_same_range(node, child) then
+      return node_normalize_down(child)
+    end
+  end
+
+  return node
+end
+
+local function visual_select(range)
+  assert(type(range) == 'table')
+  local srow, scol, erow, ecol = Range.unpack4(range)
+  local cursor_other_end_of_visual = false
+
+  if vim.fn.mode() == 'v' then
+    local vcol, vrow = vim.fn.col('v'), vim.fn.line('v')
+    local ccol, cline = vim.fn.col('.'), vim.fn.line('.')
+    if vrow > cline or (vrow == cline and vcol > ccol) then
+      cursor_other_end_of_visual = true
+    end
+  end
+
+  vim.api.nvim_win_set_cursor(0, { srow + 1, scol })
+  vim.api.nvim_feedkeys(vim.keycode('<C-\\><C-n>v'), 'nx', true)
+
+  if not pcall(vim.api.nvim_win_set_cursor, 0, { erow + 1, ecol - 1 }) then
+    vim.api.nvim_win_set_cursor(0, { erow, #vim.fn.getline(erow) })
+  end
+
+  if cursor_other_end_of_visual then
+    vim.api.nvim_feedkeys('o', 'nx', true)
+  end
+end
+
+--- @return Range4
+local function get_selection()
+  local pos1 = vim.fn.getpos('v')
+  local pos2 = vim.fn.getpos('.')
+  if pos1[2] > pos2[2] or (pos1[2] == pos2[2] and pos1[3] > pos2[3]) then
+    --- @type Range4,Range4
+    pos1, pos2 = pos2, pos1
+  end
+  local range = { pos1[2] - 1, pos1[3] - 1, pos2[2] - 1, pos2[3] }
+
+  if range[4] == #vim.fn.getline(range[3] + 1) + 1 then
+    range[3] = range[3] + 1
+    range[4] = 0
+  end
+
+  return range
+end
+
+local function get_parent_from_range(range)
+  local node, parent_chain = get_node(range)
+
+  if node == false then
+    return (assert(parent_chain[1]))
+  end
+
+  if not node then
+    return
+  end
+
+  if not range_is_same(range, node_range(node)) then
+    return node
+  end
+
+  node = node_normalize_up(node, parent_chain)
+
+  local parent = node_get_parent_no_normalize(node, parent_chain)
+
+  if parent then
+    if
+      history.bufnr ~= vim.api.nvim_get_current_buf()
+      or history.changedtick ~= vim.b.changedtick
+      or history.current_node_id ~= node_id(node)
+    then
+      history = {
+        bufnr = vim.api.nvim_get_current_buf(),
+        changedtick = vim.b.changedtick,
+      }
+    end
+    table.insert(history, node)
+    history.current_node_id = node_id(parent)
+
+    return parent
+  end
+end
+
+local function get_child_from_range(range)
+  local node, alternative_child_nodes = get_node(range)
+
+  if node == false then
+    return (assert(alternative_child_nodes[1]))
+  end
+
+  if not node then
+    return
+  end
+
+  node = node_normalize_down(node)
+
+  if not range_is_same(range, node_range(node)) then
+    history = {}
+
+    local smallest_node = get_node_contained_in_range(range, node)
+    if smallest_node then
+      return smallest_node
+    end
+
+    return node
+  end
+
+  if
+    history.bufnr == vim.api.nvim_get_current_buf()
+    and history.changedtick == vim.b.changedtick
+    and history.current_node_id == node_id(node)
+  then
+    --- @type vim.treesitter.select.node
+    local child = table.remove(history)
+    if child then
+      history.current_node_id = node_id(child)
+
+      return child
+    end
+  end
+  history = {}
+
+  for _, child in ipairs(node_get_children_no_normalize(node)) do
+    if not node_is_size_0(child) then
+      return child
+    end
+  end
+
+  return node
+end
+
+--- @param prev boolean
+local function get_sibling_from_range(range, prev)
+  local node, parent_chain = get_node(range)
+  if not node then
+    return
+  end
+
+  node = node_normalize_up(node, parent_chain)
+  local parent = node_get_parent_no_normalize(node, parent_chain)
+  if not parent then
+    return
+  end
+
+  local siblings = node_get_children_no_normalize(parent)
+
+  --- @type integer?
+  local idx
+  for n, child in ipairs(siblings) do
+    if node_id(child) == node_id(node) then
+      idx = n + (prev and -1 or 1)
+      break
+    end
+  end
+  assert(idx)
+
+  while siblings[idx] and node_is_size_0(siblings[idx]) do
+    idx = idx + (prev and -1 or 1)
+  end
+
+  if siblings[idx] then
+    return siblings[idx]
+  end
+end
+
+local function get_next_from_range(range)
+  return get_sibling_from_range(range, false)
+end
+
+local function get_prev_from_range(range)
+  return get_sibling_from_range(range, true)
+end
+
+--- @param count integer
+--- @param fn fun(range: Range4): vim.treesitter.select.node
+local function repeate_apply_range(count, fn)
+  local range = get_selection()
+
+  for _ = 1, count or 1 do
+    local node = fn(range)
+
+    if not node then
+      break
+    end
+
+    range = node_range(node)
+  end
+
+  if range and count ~= 0 then
+    visual_select(range)
+  end
+end
+
+--- @param count integer
+function M.select_parent(count)
+  repeate_apply_range(count, get_parent_from_range)
+end
+
+--- @param count integer
+function M.select_child(count)
+  repeate_apply_range(count, get_child_from_range)
+end
+
+--- @param count integer
+function M.select_next(count)
+  repeate_apply_range(count, get_next_from_range)
+end
+
+--- @param count integer
+function M.select_prev(count)
+  repeate_apply_range(count, get_prev_from_range)
+end
+
+return M