feat(move,swap): wrapping around sibling list

this adds support for wrapping around on parameter lists by querying the
surrounding scope

Author: mazunki

lua/nvim-treesitter-textobjects/move.lua

@@ -114,6 +114,19 @@ local function move(opts, query_strings, query_group)
     end
   end
 
+  -- precompute enclosing scope once so we can discard out-of-scope results
+  -- during wrap-around without re-walking the tree each iteration.
+  local scope ---@type Range6?
+  if opts.wrap then
+    for _, query_string in ipairs(query_strings) do
+      local current = shared.textobject_at_point(query_string, query_group, bufnr)
+      if current then
+        scope = shared.get_scope_range(bufnr, current)
+        if scope then break end
+      end
+    end
+  end
+
   for _ = 1, vim.v.count1 do
     local best_range ---@type Range6?
     local best_score ---@type integer
@@ -132,6 +145,12 @@ local function move(opts, query_strings, query_group)
           end
         )
 
+        if current_range then
+          if scope and not ts_range.contains(scope, current_range) then
+            current_range = nil
+          end
+        end
+
         if current_range then
           local score = scoring_function(start_, current_range)
           if not best_range then
@@ -147,6 +166,19 @@ local function move(opts, query_strings, query_group)
         end
       end
     end
+
+    -- no direct successor/precursor was found => pick first/last of siblings
+    if not best_range and scope then
+      for _, query_string in ipairs(query_strings) do
+        local siblings = shared.ranges_in_scope(bufnr, query_string, query_group, scope)
+        if #siblings > 0 then
+          best_range = forward and siblings[1] or siblings[#siblings]
+          best_start = starts[1]
+          break
+        end
+      end
+    end
+
     goto_node(best_range and best_range, not best_start, not config.set_jumps)
   end
 end
@@ -156,51 +188,61 @@ local move_repeatable = repeatable_move.make_repeatable_move(move)
 
 ---@param query_strings string|string[]
 ---@param query_group? string
-M.goto_next_start = function(query_strings, query_group)
+---@param opts? {wrap?: boolean}
+M.goto_next_start = function(query_strings, query_group, opts)
   move_repeatable({
     forward = true,
     start = true,
+    wrap = opts and opts.wrap,
   }, query_strings, query_group)
 end
 ---@param query_strings string|string[]
 ---@param query_group? string
-M.goto_next_end = function(query_strings, query_group)
+---@param opts? {wrap?: boolean}
+M.goto_next_end = function(query_strings, query_group, opts)
   move_repeatable({
     forward = true,
     start = false,
+    wrap = opts and opts.wrap,
   }, query_strings, query_group)
 end
 ---@param query_strings string|string[]
 ---@param query_group? string
-M.goto_previous_start = function(query_strings, query_group)
+---@param opts? {wrap?: boolean}
+M.goto_previous_start = function(query_strings, query_group, opts)
   move_repeatable({
     forward = false,
     start = true,
+    wrap = opts and opts.wrap,
   }, query_strings, query_group)
 end
 ---@param query_strings string|string[]
 ---@param query_group? string
-M.goto_previous_end = function(query_strings, query_group)
+---@param opts? {wrap?: boolean}
+M.goto_previous_end = function(query_strings, query_group, opts)
   move_repeatable({
     forward = false,
     start = false,
+    wrap = opts and opts.wrap,
   }, query_strings, query_group)
 end
 
 ---@param query_strings string|string[]
 ---@param query_group? string
-M.goto_next = function(query_strings, query_group)
+---@param opts? {wrap?: boolean}
+M.goto_next = function(query_strings, query_group, opts)
   move_repeatable({
     forward = true,
+    wrap = opts and opts.wrap,
   }, query_strings, query_group)
 end
 ---@param query_strings string|string[]
 ---@param query_group? string
-M.goto_previous = function(query_strings, query_group)
+---@param opts? {wrap?: boolean}
+M.goto_previous = function(query_strings, query_group, opts)
   move_repeatable({
     forward = false,
-    query_strings = query_strings,
-    query_group = query_group,
+    wrap = opts and opts.wrap,
   }, query_strings, query_group)
 end
 

lua/nvim-treesitter-textobjects/repeatable_move.lua

@@ -6,6 +6,7 @@ local M = {}
 ---@class TSTextObjects.MoveOpts
 ---@field forward boolean If true, move forward, and false is for backward.
 ---@field start? boolean If true, choose the start of the node, and false is for the end.
+---@field wrap? boolean If true, wrap around within the enclosing scope (e.g. argument list) when reaching the first or last sibling.
 
 ---@alias TSTextObjects.MoveFunction fun(opts: TSTextObjects.MoveOpts, ...: any)
 

lua/nvim-treesitter-textobjects/shared.lua

@@ -212,6 +212,50 @@ function M.find_best_range(bufnr, capture_string, query_group, filter_predicate,
   return best
 end
 
+---Get the Range6 of the treesitter parent node that encloses the given range.
+---Used to derive the "scope" (e.g. argument list) for a textobject capture.
+---@param bufnr integer
+---@param range Range6
+---@return Range6?
+function M.get_scope_range(bufnr, range)
+  local srow, scol, erow, ecol = range[1], range[2], range[4], range[5]
+  local node = ts.get_node { bufnr = bufnr, pos = { srow, scol } }
+  if not node then return nil end
+
+  -- walk up until we find a node that fully contains the textobject range,
+  -- then return its parent's range as the scope
+  while node do
+    local nsr, nsc, ner, nec = node:range()
+    if ts_range.cmp_pos.le(nsr, nsc, srow, scol) and ts_range.cmp_pos.ge(ner, nec, erow, ecol) then
+      local parent = node:parent()
+      if not parent then return nil end
+      local psr, psc, per, pec = parent:range()
+      return ts_range.add_bytes(bufnr, { psr, psc, per, pec })
+    end
+    node = node:parent()
+  end
+  return nil
+end
+
+---Get all capture ranges for `capture_string` that fall within `scope_range`,
+---sorted by start byte. Used to enumerate siblings for wrap-around navigation.
+---@param bufnr integer
+---@param capture_string string
+---@param query_group string
+---@param scope_range Range6
+---@return Range6[]
+function M.ranges_in_scope(bufnr, capture_string, query_group, scope_range)
+  local ranges = get_capture_ranges_recursively(bufnr, capture_string, query_group)
+  local result = {} ---@type Range6[]
+  for _, r in ipairs(ranges) do
+    if ts_range.contains(scope_range, r) then
+      result[#result + 1] = r
+    end
+  end
+  table.sort(result, function(a, b) return a[3] < b[3] end)
+  return result
+end
+
 -- TODO: replace with `vim.Range:has(vim.Pos)` when we drop support for nvim 0.11
 ---@param range Range
 ---@param line integer

lua/nvim-treesitter-textobjects/swap.lua

@@ -167,7 +167,8 @@ local M = {}
 ---@param query_strings string|string[]
 ---@param query_group? string
 ---@param direction integer
-local function swap_textobject(query_strings, query_group, direction)
+---@param opts? {wrap?: boolean}
+local function swap_textobject(query_strings, query_group, direction, opts)
   if type(query_strings) == 'string' then
     query_strings = { query_strings }
   end
@@ -186,11 +187,43 @@ local function swap_textobject(query_strings, query_group, direction)
     return
   end
 
+  local scope = opts and opts.wrap and shared.get_scope_range(bufnr, textobject_range) or nil
+
   local step = direction > 0 and 1 or -1
   for _ = 1, math.abs(direction), step do
     local adjacent = direction > 0
         and next_textobject(textobject_range, query_string, query_group, bufnr)
       or previous_textobject(textobject_range, query_string, query_group, bufnr)
+
+    -- next/previous_textobject searched the whole file: if the found node is
+    -- outside our scope it is not a sibling, so discard it.
+    if adjacent and scope and not ts_range.contains(scope, adjacent) then
+      adjacent = nil
+    end
+
+    -- if no adjacent sibling found and wrap is requested, find the first/last
+    -- sibling within the enclosing scope and swap with that instead.
+    if not adjacent and scope then
+      local siblings = shared.ranges_in_scope(bufnr, query_string, query_group, scope)
+      if #siblings >= 2 then
+        if direction > 0 then
+          for _, sibling in ipairs(siblings) do
+            if not range_eq(sibling, textobject_range) then
+              adjacent = sibling
+              break
+            end
+          end
+        else
+          for i = #siblings, 1, -1 do
+            if not range_eq(siblings[i], textobject_range) then
+              adjacent = siblings[i]
+              break
+            end
+          end
+        end
+      end
+    end
+
     if adjacent then
       swap_nodes(textobject_range, adjacent, bufnr, 'yes, set cursor!')
     end
@@ -206,17 +239,19 @@ end
 
 ---@param query_strings string lua pattern describing the query string
 ---@param query_group? string
-function M.swap_next(query_strings, query_group)
+---@param opts? {wrap?: boolean}
+function M.swap_next(query_strings, query_group, opts)
   return make_dot_repeatable(function()
-    swap_textobject(query_strings, query_group, 1)
+    swap_textobject(query_strings, query_group, 1, opts)
   end)
 end
 
 ---@param query_strings string lua pattern describing the query string
 ---@param query_group? string
-function M.swap_previous(query_strings, query_group)
+---@param opts? {wrap?: boolean}
+function M.swap_previous(query_strings, query_group, opts)
   return make_dot_repeatable(function()
-    swap_textobject(query_strings, query_group, -1)
+    swap_textobject(query_strings, query_group, -1, opts)
   end)
 end