feat(treesitter): add 'injection.self' and 'injection.parent'

Co-authored-by: ObserverOfTime <chronobserver@disroot.org>

4 files changed, 47 insertions, 3 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 483e829770..82b390853c 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -131,6 +131,8 @@ The following new APIs and features were added.
   • `@injection.language` now has smarter resolution and will now fallback to language aliases and/or attempt lower case variants of the text.
     language via aliases (e.g., filetype) registered via
     `vim.treesitter.language.register`.
+  • The `#set!` directive now supports `injection.self` and `injection.parent` for injecting either the current node's language
+    or the parent LanguageTree's language, respectively.
 
 • |vim.ui.open()| opens URIs using the system default handler (macOS `open`,
   Windows `explorer`, Linux `xdg-open`, etc.)
diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index f3e697807f..139b986786 100644
--- a/runtime/doc/treesitter.txt
+++ b/runtime/doc/treesitter.txt
@@ -523,6 +523,10 @@ associated with patterns:
     node's entire text should be re-parsed, including the text of its child
     nodes. By default, child nodes' text will be excluded from the injected
     document.
+    • `injection.self` - indicates that the node's text should be parsed with
+      the same language as the node's LanguageTree.
+    • `injection.parent` - indicates that the captured node’s text should
+      be parsed with the same language as the node's parent LanguageTree.
 
 ==============================================================================
 VIM.TREESITTER                                                *lua-treesitter*
diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua
index 4b419c4744..e81778b269 100644
--- a/runtime/lua/vim/treesitter/languagetree.lua
+++ b/runtime/lua/vim/treesitter/languagetree.lua
@@ -82,6 +82,7 @@ local TSCallbackNames = {
 ---List of regions this tree should manage and parse. If nil then regions are
 ---taken from _trees. This is mostly a short-lived cache for included_regions()
 ---@field private _lang string Language name
+---@field private _parent_lang? string Parent language name
 ---@field private _source (integer|string) Buffer or string to parse
 ---@field private _trees TSTree[] Reference to parsed tree (one for each language)
 ---@field private _valid boolean|table<integer,boolean> If the parsed tree is valid
@@ -105,8 +106,9 @@ LanguageTree.__index = LanguageTree
 ---@param opts (table|nil) Optional arguments:
 ---             - injections table Map of language to injection query strings. Overrides the
 ---                                built-in runtime file searching for language injections.
+---@param parent_lang? string Parent language name of this tree
 ---@return LanguageTree parser object
-function LanguageTree.new(source, lang, opts)
+function LanguageTree.new(source, lang, opts, parent_lang)
   language.add(lang)
   ---@type LanguageTreeOpts
   opts = opts or {}
@@ -121,6 +123,7 @@ function LanguageTree.new(source, lang, opts)
   local self = {
     _source = source,
     _lang = lang,
+    _parent_lang = parent_lang,
     _children = {},
     _trees = {},
     _opts = opts,
@@ -489,7 +492,7 @@ function LanguageTree:add_child(lang)
     self:remove_child(lang)
   end
 
-  local child = LanguageTree.new(self._source, lang, self._opts)
+  local child = LanguageTree.new(self._source, lang, self._opts, self:lang())
 
   -- Inherit recursive callbacks
   for nm, cb in pairs(self._callbacks_rec) do
@@ -752,7 +755,9 @@ end
 function LanguageTree:_get_injection(match, metadata)
   local ranges = {} ---@type Range6[]
   local combined = metadata['injection.combined'] ~= nil
-  local lang = metadata['injection.language'] --[[@as string?]]
+  local lang = metadata['injection.self'] ~= nil and self:lang()
+    or metadata['injection.parent'] ~= nil and self._parent_lang
+    or metadata['injection.language'] --[[@as string?]]
   local include_children = metadata['injection.include-children'] ~= nil
 
   for id, node in pairs(match) do
diff --git a/test/functional/treesitter/parser_spec.lua b/test/functional/treesitter/parser_spec.lua
index ae7cf58293..56af0c8738 100644
--- a/test/functional/treesitter/parser_spec.lua
+++ b/test/functional/treesitter/parser_spec.lua
@@ -696,6 +696,39 @@ int x = INT_MAX;
       end)
     end)
 
+    describe("when using injection.self", function()
+      it("should inject the source language", function()
+        exec_lua([[
+        parser = vim.treesitter.get_parser(0, "c", {
+          injections = {
+            c = '(preproc_def (preproc_arg) @injection.content (#set! injection.self)) (preproc_function_def value: (preproc_arg) @injection.content (#set! injection.self))'}})
+        parser:parse(true)
+        ]])
+
+        eq("table", exec_lua("return type(parser:children().c)"))
+        eq(5, exec_lua("return #parser:children().c:trees()"))
+        eq({
+          {0, 0, 7, 0},   -- root tree
+          {3, 14, 3, 17}, -- VALUE 123
+          {4, 15, 4, 18}, -- VALUE1 123
+          {5, 15, 5, 18}, -- VALUE2 123
+          {1, 26, 1, 63}, -- READ_STRING(x, y) (char *)read_string((x), (size_t)(y))
+          {2, 29, 2, 66}  -- READ_STRING_OK(x, y) (char *)read_string((x), (size_t)(y))
+        }, get_ranges())
+
+        helpers.feed('ggo<esc>')
+        eq(5, exec_lua("return #parser:children().c:trees()"))
+        eq({
+          {0, 0, 8, 0},   -- root tree
+          {4, 14, 4, 17}, -- VALUE 123
+          {5, 15, 5, 18}, -- VALUE1 123
+          {6, 15, 6, 18}, -- VALUE2 123
+          {2, 26, 2, 63}, -- READ_STRING(x, y) (char *)read_string((x), (size_t)(y))
+          {3, 29, 3, 66}  -- READ_STRING_OK(x, y) (char *)read_string((x), (size_t)(y))
+        }, get_ranges())
+      end)
+    end)
+
     describe("when using the offset directive", function()
       it("should shift the range by the directive amount", function()
         exec_lua([[