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

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

3 files changed, 14 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