feat(treesitter)!: apply offset! directive to all captures #34383

This commit changes the `offset!` directive so that instead of setting a
`metadata.range` value for the entire pattern, it will set a
`metadata.offset` value. This offset will be applied to the range only
in `vim.treesitter.get_range()`, rather than at directive application
time. This allows the offset to be applied to any and all nodes captured
by the given pattern, and removes the requirement that `#offset!` be
applied to only a single node.

The downside of this change is that plugins which read from
`metadata.range` may be thrown off course, but such plugins should
prefer `vim.treesitter.get_range()` when retrieving ranges anyway.

Note that `#trim!` still sets `metadata.range`, and
`vim.treesitter.get_range()` still reads from `metadata.range`, if it
exists.

runtime/lua/vim/treesitter.lua

@@ -165,6 +165,31 @@ function M.get_node_range(node_or_range)
   end
 end
 
+---@param node TSNode
+---@param source integer|string Buffer or string from which the {node} is extracted
+---@param offset Range4
+---@return Range6
+local function apply_range_offset(node, source, offset)
+  ---@diagnostic disable-next-line: missing-fields LuaLS varargs bug
+  local range = { node:range() } ---@type Range4
+  local start_row_offset = offset[1]
+  local start_col_offset = offset[2]
+  local end_row_offset = offset[3]
+  local end_col_offset = offset[4]
+
+  range[1] = range[1] + start_row_offset
+  range[2] = range[2] + start_col_offset
+  range[3] = range[3] + end_row_offset
+  range[4] = range[4] + end_col_offset
+
+  if range[1] < range[3] or (range[1] == range[3] and range[2] <= range[4]) then
+    return M._range.add_bytes(source, range)
+  end
+
+  -- If this produces an invalid range, we just skip it.
+  return { node:range(true) }
+end
+
 ---Get the range of a |TSNode|. Can also supply {source} and {metadata}
 ---to get the range with directives applied.
 ---@param node TSNode
@@ -172,9 +197,12 @@ end
 ---@param metadata vim.treesitter.query.TSMetadata|nil
 ---@return Range6
 function M.get_range(node, source, metadata)
-  if metadata and metadata.range then
-    assert(source)
-    return M._range.add_bytes(source, metadata.range)
+  if metadata then
+    if metadata.range then
+      return M._range.add_bytes(assert(source), metadata.range)
+    elseif metadata.offset then
+      return apply_range_offset(node, assert(source), metadata.offset)
+    end
   end
   return { node:range(true) }
 end

runtime/lua/vim/treesitter/query.lua

@@ -607,6 +607,7 @@ predicate_handlers['any-vim-match?'] = predicate_handlers['any-match?']
 ---@nodoc
 ---@class vim.treesitter.query.TSMetadata
 ---@field range? Range
+---@field offset? Range4
 ---@field conceal? string
 ---@field bo.commentstring? string
 ---@field [integer]? vim.treesitter.query.TSMetadata
@@ -645,29 +646,21 @@ local directive_handlers = {
     if not nodes or #nodes == 0 then
       return
     end
-    assert(#nodes == 1, '#offset! does not support captures on multiple nodes')
-
-    local node = nodes[1]
 
     if not metadata[capture_id] then
       metadata[capture_id] = {}
     end
 
-    local range = metadata[capture_id].range or { node:range() }
-    local start_row_offset = pred[3] or 0
-    local start_col_offset = pred[4] or 0
-    local end_row_offset = pred[5] or 0
-    local end_col_offset = pred[6] or 0
-
-    range[1] = range[1] + start_row_offset
-    range[2] = range[2] + start_col_offset
-    range[3] = range[3] + end_row_offset
-    range[4] = range[4] + end_col_offset
-
-    -- If this produces an invalid range, we just skip it.
-    if range[1] < range[3] or (range[1] == range[3] and range[2] <= range[4]) then
-      metadata[capture_id].range = range
-    end
+    metadata[capture_id].offset = {
+      pred[3] --[[@as integer]]
+        or 0,
+      pred[4] --[[@as integer]]
+        or 0,
+      pred[5] --[[@as integer]]
+        or 0,
+      pred[6] --[[@as integer]]
+        or 0,
+    }
   end,
   -- Transform the content of the node
   -- Example: (#gsub! @_node ".*%.(.*)" "%1")