perf(treesitter): only search for injections within the parse range

Co-authored-by: Jaehwang Jung <tomtomjhj@gmail.com>

runtime/doc/news.txt

@@ -340,6 +340,9 @@ PERFORMANCE
 • Treesitter highlighting is now asynchronous. To force synchronous parsing,
   use `vim.g._ts_force_sync_parsing = true`.
 • Treesitter folding is now calculated asynchronously.
+• |LanguageTree:parse()| now only runs the injection query on the provided
+  range (as long as the language does not have a combined injection),
+  significantly improving |treesitter-highlight| performance.
 
 PLUGINS
 

runtime/lua/vim/treesitter/languagetree.lua

@@ -46,6 +46,9 @@ local Range = require('vim.treesitter._range')
 
 local default_parse_timeout_ms = 3
 
+---@type Range2
+local entire_document_range = { 0, math.huge }
+
 ---@alias TSCallbackName
 ---| 'changedtree'
 ---| 'bytes'
@@ -77,7 +80,7 @@ local TSCallbackNames = {
 ---@field package _callbacks_rec table<TSCallbackName,function[]> Callback handlers (recursive)
 ---@field private _children table<string,vim.treesitter.LanguageTree> Injected languages
 ---@field private _injection_query vim.treesitter.Query Queries defining injected languages
----@field private _injections_processed boolean
+---@field private _processed_injection_range Range? Range for which injections have been processed
 ---@field private _opts table Options
 ---@field private _parser TSParser Parser for language
 ---Table of regions for which the tree is currently running an async parse
@@ -137,7 +140,7 @@ function LanguageTree.new(source, lang, opts)
     _opts = opts,
     _injection_query = injections[lang] and query.parse(lang, injections[lang])
       or query.get(lang, 'injections'),
-    _injections_processed = false,
+    _processed_injection_range = nil,
     _valid_regions = {},
     _num_valid_regions = 0,
     _num_regions = 1,
@@ -334,7 +337,10 @@ function LanguageTree:is_valid(exclude_children, range)
   end
 
   if not exclude_children then
-    if not self._injections_processed then
+    if
+      not self._processed_injection_range
+      or not Range.contains(self._processed_injection_range, range or entire_document_range)
+    then
       return false
     end
 
@@ -416,11 +422,12 @@ function LanguageTree:_parse_regions(range, thread_state)
 end
 
 --- @private
+--- @param range Range|true
 --- @return number
-function LanguageTree:_add_injections()
+function LanguageTree:_add_injections(range)
   local seen_langs = {} ---@type table<string,boolean>
 
-  local query_time, injections_by_lang = tcall(self._get_injections, self)
+  local query_time, injections_by_lang = tcall(self._get_injections, self, range)
   for lang, injection_regions in pairs(injections_by_lang) do
     local has_lang = pcall(language.add, lang)
 
@@ -604,13 +611,21 @@ function LanguageTree:_parse(range, thread_state)
     end
     -- Need to run injections when we parsed something
     if no_regions_parsed > 0 then
-      self._injections_processed = false
+      self._processed_injection_range = nil
     end
   end
 
-  if not self._injections_processed and range then
-    query_time = self:_add_injections()
-    self._injections_processed = true
+  if
+    range
+    and not (
+      self._processed_injection_range
+      and Range.contains(
+        self._processed_injection_range,
+        range ~= true and range or entire_document_range
+      )
+    )
+  then
+    query_time = self:_add_injections(range)
   end
 
   self:_log({
@@ -986,18 +1001,27 @@ end
 --- TODO: Allow for an offset predicate to tailor the injection range
 ---       instead of using the entire nodes range.
 --- @private
+--- @param range Range|true
 --- @return table<string, Range6[][]>
-function LanguageTree:_get_injections()
+function LanguageTree:_get_injections(range)
   if not self._injection_query or #self._injection_query.captures == 0 then
+    self._processed_injection_range = entire_document_range
     return {}
   end
 
   ---@type table<integer,vim.treesitter.languagetree.Injection>
   local injections = {}
 
+  local full_scan = range == true or self._injection_query.has_combined_injections
+
   for index, tree in pairs(self._trees) do
     local root_node = tree:root()
-    local start_line, _, end_line, _ = root_node:range()
+    local start_line, end_line ---@type integer, integer
+    if full_scan then
+      start_line, _, end_line = root_node:range()
+    else
+      start_line, _, end_line = Range.unpack4(range --[[@as Range]])
+    end
 
     for pattern, match, metadata in
       self._injection_query:iter_matches(root_node, self._source, start_line, end_line + 1)
@@ -1034,6 +1058,12 @@ function LanguageTree:_get_injections()
     end
   end
 
+  if full_scan then
+    self._processed_injection_range = entire_document_range
+  else
+    self._processed_injection_range = range --[[@as Range]]
+  end
+
   return result
 end
 

runtime/lua/vim/treesitter/query.lua

@@ -30,9 +30,11 @@ end
 --- Splits the query patterns into predicates and directives.
 ---@param patterns table<integer, (integer|string)[][]>
 ---@return table<integer, vim.treesitter.query.ProcessedPattern>
+---@return boolean
 local function process_patterns(patterns)
   ---@type table<integer, vim.treesitter.query.ProcessedPattern>
   local processed_patterns = {}
+  local has_combined = false
 
   for k, pattern_list in pairs(patterns) do
     ---@type vim.treesitter.query.ProcessedPredicate[]
@@ -47,6 +49,9 @@ local function process_patterns(patterns)
 
       if is_directive(pred_name) then
         table.insert(directives, pattern)
+        if vim.deep_equal(pattern, { 'set!', 'injection.combined' }) then
+          has_combined = true
+        end
       else
         local should_match = true
         if pred_name:match('^not%-') then
@@ -60,7 +65,7 @@ local function process_patterns(patterns)
     processed_patterns[k] = { predicates = predicates, directives = directives }
   end
 
-  return processed_patterns
+  return processed_patterns, has_combined
 end
 
 ---@nodoc
@@ -71,6 +76,7 @@ end
 ---@field captures string[] list of (unique) capture names defined in query
 ---@field info vim.treesitter.QueryInfo query context (e.g. captures, predicates, directives)
 ---@field query TSQuery userdata query object
+---@field has_combined_injections boolean whether the query contains combined injections
 ---@field private _processed_patterns table<integer, vim.treesitter.query.ProcessedPattern>
 local Query = {}
 Query.__index = Query
@@ -90,7 +96,7 @@ function Query.new(lang, ts_query)
     patterns = query_info.patterns,
   }
   self.captures = self.info.captures
-  self._processed_patterns = process_patterns(self.info.patterns)
+  self._processed_patterns, self.has_combined_injections = process_patterns(self.info.patterns)
   return self
 end
 

test/functional/treesitter/parser_spec.lua

@@ -633,7 +633,7 @@ int x = INT_MAX;
         }, get_ranges())
 
         n.feed('7ggI//<esc>')
-        exec_lua([[parser:parse({5, 6})]])
+        exec_lua([[parser:parse(true)]])
         eq('table', exec_lua('return type(parser:children().c)'))
         eq(2, exec_lua('return #parser:children().c:trees()'))
         eq({
@@ -1122,7 +1122,7 @@ print()
     )
 
     eq(
-      2,
+      1,
       exec_lua(function()
         _G.parser:parse({ 2, 6 })
         return #_G.parser:children().lua:trees()
@@ -1172,10 +1172,10 @@ print()
       eq(true, exec_lua('return vim.treesitter.get_parser():is_valid()'))
     end)
 
-    it('is fully valid after a parsing a range on parsed tree', function()
+    it('is valid within a range on parsed tree after parsing it', function()
       exec_lua('vim.treesitter.get_parser():parse({5, 7})')
       eq(true, exec_lua('return vim.treesitter.get_parser():is_valid(true)'))
-      eq(true, exec_lua('return vim.treesitter.get_parser():is_valid()'))
+      eq(true, exec_lua('return vim.treesitter.get_parser():is_valid(nil, {5, 7})'))
     end)
 
     describe('when adding content with injections', function()
@@ -1200,14 +1200,11 @@ print()
         eq(false, exec_lua('return vim.treesitter.get_parser():is_valid()'))
       end)
 
-      it(
-        'is fully valid after a range parse that leads to parsing not parsed injections',
-        function()
-          exec_lua('vim.treesitter.get_parser():parse({5, 7})')
-          eq(true, exec_lua('return vim.treesitter.get_parser():is_valid(true)'))
-          eq(true, exec_lua('return vim.treesitter.get_parser():is_valid()'))
-        end
-      )
+      it('is valid within a range on parsed tree after parsing it', function()
+        exec_lua('vim.treesitter.get_parser():parse({5, 7})')
+        eq(true, exec_lua('return vim.treesitter.get_parser():is_valid(true)'))
+        eq(true, exec_lua('return vim.treesitter.get_parser():is_valid(nil, {5, 7})'))
+      end)
 
       it(
         'is valid excluding, invalid including children after a range parse that does not lead to parsing not parsed injections',
@@ -1249,10 +1246,10 @@ print()
         eq(false, exec_lua('return vim.treesitter.get_parser():is_valid()'))
       end)
 
-      it('is fully valid after a range parse that leads to parsing modified child tree', function()
+      it('is valid within a range parse that leads to parsing modified child tree', function()
         exec_lua('vim.treesitter.get_parser():parse({5, 7})')
         eq(true, exec_lua('return vim.treesitter.get_parser():is_valid(true)'))
-        eq(true, exec_lua('return vim.treesitter.get_parser():is_valid()'))
+        eq(true, exec_lua('return vim.treesitter.get_parser():is_valid(nil, {5, 7})'))
       end)
 
       it(