1 files changed, 74 insertions, 14 deletions

diff --git a/runtime/lua/vim/treesitter/_select.lua b/runtime/lua/vim/treesitter/_select.lua
index ffca1a1ce5..c3986b486d 100644
--- a/runtime/lua/vim/treesitter/_select.lua
+++ b/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