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/doc/news.txt

@@ -88,7 +88,11 @@ PLUGINS
 
 TREESITTER
 
-• todo
+• |treesitter-directive-offset!| can now be applied to quantified captures. It
+  no longer sets `metadata[capture_id].range`; it instead sets
+  `metadata[capture_id].offset`. The offset will be applied in
+  |vim.treesitter.get_range()|, which should be preferred over reading
+  metadata directly for retrieving node ranges.
 
 TUI
 

runtime/doc/treesitter.txt

@@ -219,7 +219,9 @@ The following directives are built in:
         Takes the range of the captured node and applies an offset. This will
         set a new range in the form of a list like { {start_row}, {start_col},
         {end_row}, {end_col} } for the captured node with `capture_id` as
-        `metadata[capture_id].range`. Useful for |treesitter-language-injections|.
+        `metadata[capture_id].offset`. This offset will be applied to the
+        range returned in |vim.treesitter.get_range()|. Useful for
+        |treesitter-language-injections|.
 
         Parameters: ~
             {capture_id}

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")

test/functional/treesitter/parser_spec.lua

@@ -838,6 +838,38 @@ int x = INT_MAX;
           { 5, 17, 5, 17 }, -- VALUE2 123
         }, get_ranges())
       end)
+      it('should apply offsets to quantified captures', function()
+        local function get_ltree_ranges()
+          return exec_lua(function()
+            local result = {}
+            _G.parser:for_each_tree(function(_, ltree)
+              table.insert(result, ltree:included_regions())
+            end)
+            return result
+          end)
+        end
+
+        exec_lua(function()
+          _G.parser = vim.treesitter.get_parser(0, 'c', {
+            injections = {
+              c = '((preproc_def (preproc_arg) @injection.content)+ (#set! injection.language "c") (#offset! @injection.content 0 1 0 -1))',
+            },
+          })
+          _G.parser:parse(true)
+        end)
+
+        eq('table', exec_lua('return type(parser:children().c)'))
+        eq({
+          { {} }, -- root tree
+          {
+            {
+              { 3, 15, 163, 3, 16, 164 }, -- VALUE 123
+              { 4, 16, 182, 4, 17, 183 }, -- VALUE1 123
+              { 5, 16, 201, 5, 17, 202 }, -- VALUE2 123
+            },
+          },
+        }, get_ltree_ranges())
+      end)
       it('should list all directives', function()
         local res_list = exec_lua(function()
           local query = vim.treesitter.query

test/functional/treesitter/testutil.lua

@@ -11,12 +11,13 @@ function M.run_query(language, query_string)
     local query = vim.treesitter.query.parse(lang, query_str)
     local parser = vim.treesitter.get_parser()
     local tree = parser:parse()[1]
+    local Range = require('vim.treesitter._range')
     local res = {}
     for id, node, metadata in query:iter_captures(tree:root(), 0) do
-      table.insert(
-        res,
-        { query.captures[id], metadata[id] and metadata[id].range or { node:range() } }
-      )
+      table.insert(res, {
+        query.captures[id],
+        { Range.unpack4(vim.treesitter.get_range(node, 0, metadata[id])) },
+      })
     end
     return res
   end, language, query_string)