perf: much faster treesitter highlighter

lua/trouble/view/text.lua

@@ -138,7 +138,7 @@ function M:render(buf)
   vim.api.nvim_buf_clear_namespace(buf, M.ns, 0, -1)
   vim.api.nvim_buf_set_lines(buf, 0, -1, false, lines)
 
-  local regions = {} ---@type table<string, number[][][]>
+  local regions = {} ---@type trouble.LangRegions
 
   for l, line in ipairs(self._lines) do
     local col = self._opts.padding or 0
@@ -174,12 +174,15 @@ function M:render(buf)
       col = col + width
     end
   end
-
-  for lang, r in pairs(regions) do
-    local TS = require("trouble.view.treesitter")
-    TS.highlight(buf, lang, r)
-  end
   vim.bo[buf].modifiable = false
+  local changetick = vim.b[buf].changetick
+
+  vim.schedule(function()
+    if vim.b[buf].changetick ~= changetick then
+      return
+    end
+    require("trouble.view.treesitter").attach(buf, regions)
+  end)
 end
 
 function M:trim()

lua/trouble/view/treesitter.lua

@@ -1,64 +1,89 @@
-local Util = require("trouble.util")
+---@alias trouble.LangRegions table<string, number[][][]>
 
 local M = {}
 
----@param buf number
----@param lang? string
-function M.highlight(buf, lang, regions)
-  local Render = require("trouble.view.render")
-  lang = lang or "markdown"
-  lang = lang == "markdown" and "markdown_inline" or lang
-  -- lang = "markdown_inline"
-  local ok, parser = pcall(vim.treesitter.get_parser, buf, lang)
+M.cache = {} ---@type table<number, table<string,{parser: vim.treesitter.LanguageTree, highlighter:vim.treesitter.highlighter, enabled:boolean}>>
+local ns = vim.api.nvim_create_namespace("trouble.treesitter")
+
+local TSHighlighter = vim.treesitter.highlighter
+
+local function wrap(name)
+  return function(_, win, buf, ...)
+    if not M.cache[buf] then
+      return false
+    end
+    for _, hl in pairs(M.cache[buf] or {}) do
+      if hl.enabled then
+        TSHighlighter.active[buf] = hl.highlighter
+        TSHighlighter[name](_, win, buf, ...)
+      end
+    end
+    TSHighlighter.active[buf] = nil
+  end
+end
 
-  if not ok then
-    local msg = "nvim-treesitter parser missing `" .. lang .. "`"
-    vim.notify_once(msg, vim.log.levels.WARN, { title = "trouble.nvim" })
+M.did_setup = false
+function M.setup()
+  if M.did_setup then
     return
   end
+  M.did_setup = true
 
-  ---@diagnostic disable-next-line: invisible
-  parser:set_included_regions(regions)
-  parser:parse(true)
-  local ret = {} ---@type Extmark[]
+  vim.api.nvim_set_decoration_provider(ns, {
+    on_win = wrap("_on_win"),
+    on_line = wrap("_on_line"),
+  })
 
-  parser:for_each_tree(function(tstree, tree)
-    if not tstree then
-      return
-    end
-    local query = vim.treesitter.query.get(tree:lang(), "highlights")
+  vim.api.nvim_create_autocmd("BufWipeout", {
+    group = vim.api.nvim_create_augroup("trouble.treesitter.hl", { clear = true }),
+    callback = function(ev)
+      M.cache[ev.buf] = nil
+    end,
+  })
+end
 
-    -- Some injected languages may not have highlight queries.
-    if not query then
-      return
-    end
+---@param buf number
+---@param regions trouble.LangRegions
+function M.attach(buf, regions)
+  M.setup()
+  M.cache[buf] = M.cache[buf] or {}
+  for lang in pairs(M.cache[buf]) do
+    M.cache[buf][lang].enabled = regions[lang] ~= nil
+  end
 
-    ---@diagnostic disable-next-line: missing-parameter
-    local iter = query:iter_captures(tstree:root(), buf)
+  for lang in pairs(regions) do
+    M._attach_lang(buf, lang, regions[lang])
+  end
+end
 
-    for capture, node, metadata in iter do
-      ---@type number, number, number, number
-      local start_row, start_col, end_row, end_col = node:range()
+---@param buf number
+---@param lang? string
+function M._attach_lang(buf, lang, regions)
+  lang = lang or "markdown"
+  lang = lang == "markdown" and "markdown_inline" or lang
 
-      ---@type string
-      local name = query.captures[capture]
-      local hl = 0
-      if not vim.startswith(name, "_") then
-        hl = vim.api.nvim_get_hl_id_by_name("@" .. name .. "." .. lang)
-      end
+  M.cache[buf] = M.cache[buf] or {}
 
-      if hl and name ~= "spell" then
-        Util.set_extmark(buf, Render.ns, start_row, start_col, {
-          end_line = end_row,
-          end_col = end_col,
-          hl_group = hl,
-          priority = (tonumber(metadata.priority) or 100) + 10, -- add 10, so it will be higher than the standard highlighter of the buffer
-          conceal = metadata.conceal,
-        })
-      end
+  if not M.cache[buf][lang] then
+    local ok, parser = pcall(vim.treesitter.get_parser, buf, lang)
+    if not ok then
+      local msg = "nvim-treesitter parser missing `" .. lang .. "`"
+      vim.notify_once(msg, vim.log.levels.WARN, { title = "trouble.nvim" })
+      return
     end
-  end)
-  return ret
+
+    parser:set_included_regions(regions)
+
+    M.cache[buf][lang] = {
+      parser = parser,
+      highlighter = TSHighlighter.new(parser),
+    }
+  end
+  M.cache[buf][lang].enabled = true
+  local parser = M.cache[buf][lang].parser
+
+  parser:set_included_regions(regions)
+  parser:parse(true)
 end
 
 return M