-
Notifications
You must be signed in to change notification settings - Fork 177
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
perf: much faster treesitter highlighter
- Loading branch information
Showing
2 changed files
with
81 additions
and
53 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |