backport feat(treesitter): expand selection to sibling node (#39323)

Problem:
Can't expand treesitter-incremental-selection to the next and previous
sibling nodes.

Solution:
Pressing `]N` in visual mode will expand the selection to the next
sibling node, and `[N` will do the same with the previous node.

Co-authored-by: altermo <107814000+altermo@users.noreply.github.com>

runtime/doc/news.txt

@@ -395,6 +395,7 @@ TREESITTER
 • |:EditQuery| command gained tab-completion, works with injected languages.
 • |LanguageTree:parse()| now accepts a list of ranges.
 • |v_an| |v_in| |v_]n| |v_[n| incremental selection of treesitter nodes.
+• |v_]N| |v_[N| expand selection to sibling treesitter node.
 
 TUI
 

runtime/doc/treesitter.txt

@@ -628,6 +628,12 @@ in                      Selects [count]th previous (or first) child node.
                                                                         *v_[n*
 [n                      Selects [count]th previous node.
 
+                                                                        *v_]N*
+]N                      Expands selection to [count]th next node.
+
+                                                                        *v_[N*
+[N                      Expands selection to [count]th previous node.
+
 ==============================================================================
 VIM.TREESITTER                                                *lua-treesitter*
 

runtime/lua/vim/_core/defaults.lua

@@ -459,6 +459,14 @@ do
       require 'vim.treesitter._select'.select_next(vim.v.count1)
     end, { desc = 'Select next node' })
 
+    vim.keymap.set({ 'x' }, '[N', function()
+      require 'vim.treesitter._select'.select_grow_prev(vim.v.count1)
+    end, { desc = 'Select expand previous node' })
+
+    vim.keymap.set({ 'x' }, ']N', function()
+      require 'vim.treesitter._select'.select_grow_next(vim.v.count1)
+    end, { desc = 'Select expand next node' })
+
     vim.keymap.set({ 'x', 'o' }, 'an', function()
       if vim.treesitter.get_parser(nil, nil, { error = false }) then
         require 'vim.treesitter._select'.select_parent(vim.v.count1)

runtime/lua/vim/treesitter/_select.lua

@@ -401,7 +401,7 @@ local function get_parent_from_range(range)
   local node, parent_chain = get_node(range)
 
   if node == false then
-    return (assert(parent_chain[1]))
+    return node_range(assert(parent_chain[1]))
   end
 
   if not node then
@@ -409,7 +409,7 @@ local function get_parent_from_range(range)
   end
 
   if not Range.equal(range, node_range(node)) then
-    return node
+    return node_range(node)
   end
 
   node = node_normalize_up(node, parent_chain)
@@ -430,7 +430,7 @@ local function get_parent_from_range(range)
     table.insert(history, node)
     history.current_node_id = node_id(parent)
 
-    return parent
+    return node_range(parent)
   end
 end
 
@@ -438,7 +438,7 @@ 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]))
+    return node_range(assert(alternative_child_nodes[1]))
   end
 
   if not node then
@@ -452,10 +452,10 @@ local function get_child_from_range(range)
 
     local smallest_node = get_node_contained_in_range(range, node)
     if smallest_node then
-      return smallest_node
+      return node_range(smallest_node)
     end
 
-    return node
+    return node_range(node)
   end
 
   if
@@ -468,18 +468,18 @@ local function get_child_from_range(range)
     if child then
       history.current_node_id = node_id(child)
 
-      return child
+      return node_range(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
+      return node_range(child)
     end
   end
 
-  return node
+  return node_range(node)
 end
 
 --- @param prev boolean
@@ -512,7 +512,7 @@ local function get_sibling_from_range(range, prev)
   end
 
   if siblings[idx] then
-    return siblings[idx]
+    return node_range(siblings[idx])
   end
 end
 
@@ -524,19 +524,69 @@ local function get_prev_from_range(range)
   return get_sibling_from_range(range, true)
 end
 
+--- @param prev boolean
+local function get_grow_sibling_from_range(range, prev)
+  local node, parent_chain = get_node(range)
+  if not node then
+    return
+  end
+
+  if Range.equal(node_range(node), range) then
+    node = node_normalize_up(node, parent_chain)
+    node = node_get_parent_no_normalize(node, parent_chain)
+
+    if not node then
+      return
+    end
+  else
+    node = node_normalize_down(node)
+  end
+
+  local children = node_get_children_no_normalize(node)
+
+  if prev then
+    for idx = #children, 1, -1 do
+      local child = children[idx]
+      local crange = node_range(child)
+      if
+        not node_is_size_0(child) and Range.cmp_pos.lt(crange[1], crange[2], range[1], range[2])
+      then
+        return { crange[1], crange[2], range[3], range[4] }
+      end
+    end
+  else
+    for _, child in ipairs(children) do
+      local crange = node_range(child)
+      if
+        not node_is_size_0(child) and Range.cmp_pos.gt(crange[3], crange[4], range[3], range[4])
+      then
+        return { range[1], range[2], crange[3], crange[4] }
+      end
+    end
+  end
+end
+
+local function get_grow_next_from_range(range)
+  return get_grow_sibling_from_range(range, false)
+end
+
+local function get_grow_prev_from_range(range)
+  return get_grow_sibling_from_range(range, true)
+end
+
 --- @param count integer
---- @param fn fun(range: Range4): vim.treesitter.select.node
+--- @param fn fun(range: Range4): Range4?
 local function repeate_apply_range(count, fn)
   local range = get_selection()
 
   for _ = 1, count or 1 do
-    local node = fn(range)
+    local new_range = fn(range)
 
-    if not node then
+    if not new_range then
       break
     end
 
-    range = node_range(node)
+    range = new_range
   end
 
   if range and count ~= 0 then
@@ -564,4 +614,14 @@ function M.select_prev(count)
   repeate_apply_range(count, get_prev_from_range)
 end
 
+--- @param count integer
+function M.select_grow_next(count)
+  repeate_apply_range(count, get_grow_next_from_range)
+end
+
+--- @param count integer
+function M.select_grow_prev(count)
+  repeate_apply_range(count, get_grow_prev_from_range)
+end
+
 return M

test/functional/treesitter/select_spec.lua

@@ -64,6 +64,16 @@ describe('treesitter incremental-selection', function()
 
     treeselect('select_parent')
     eq('foo(1)\nbar(2)\n', get_selected())
+
+    set_lines('quux(1,foo,bar,baz,qux,2)')
+    feed('<esc>fbve')
+    eq('bar', get_selected())
+
+    treeselect('select_grow_next')
+    eq('bar,baz', get_selected())
+
+    treeselect('select_grow_prev')
+    eq('foo,bar,baz', get_selected())
   end)
 
   it('repeat', function()
@@ -89,6 +99,14 @@ describe('treesitter incremental-selection', function()
 
     treeselect('select_child', 2)
     eq('2', get_selected())
+
+    feed('<esc>F1')
+    treeselect('select_grow_next', 2)
+    eq('1,2,3', get_selected())
+
+    feed('<esc>f4v')
+    treeselect('select_grow_prev', 2)
+    eq('2,3,4', get_selected())
   end)
 
   it('history', function()