highlighter.lua (18437B)
1 local api = vim.api 2 local query = vim.treesitter.query 3 local Range = require('vim.treesitter._range') 4 local cmp_lt = Range.cmp_pos.lt 5 6 local ns = api.nvim_create_namespace('nvim.treesitter.highlighter') 7 8 ---@alias vim.treesitter.highlighter.Iter fun(end_line: integer|nil, end_col: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata, TSQueryMatch, TSTree 9 10 ---@class (private) vim.treesitter.highlighter.Query 11 ---@field private _query vim.treesitter.Query? 12 ---@field private lang string 13 ---@field private hl_cache table<integer,integer> 14 local TSHighlighterQuery = {} 15 TSHighlighterQuery.__index = TSHighlighterQuery 16 17 ---@private 18 ---@param lang string 19 ---@param query_string string? 20 ---@return vim.treesitter.highlighter.Query 21 function TSHighlighterQuery.new(lang, query_string) 22 local self = setmetatable({}, TSHighlighterQuery) 23 self.lang = lang 24 self.hl_cache = {} 25 26 if query_string then 27 self._query = query.parse(lang, query_string) 28 else 29 self._query = query.get(lang, 'highlights') 30 end 31 32 return self 33 end 34 35 ---@package 36 ---@param capture integer 37 ---@return integer? 38 function TSHighlighterQuery:get_hl_from_capture(capture) 39 if not self.hl_cache[capture] then 40 local name = self._query.captures[capture] 41 local id = 0 42 if not vim.startswith(name, '_') then 43 id = api.nvim_get_hl_id_by_name('@' .. name .. '.' .. self.lang) 44 end 45 self.hl_cache[capture] = id 46 end 47 48 return self.hl_cache[capture] 49 end 50 51 ---@nodoc 52 function TSHighlighterQuery:query() 53 return self._query 54 end 55 56 ---@class (private) vim.treesitter.highlighter.State 57 ---@field tstree TSTree 58 ---@field next_row integer 59 ---@field next_col integer 60 ---@field iter vim.treesitter.highlighter.Iter? 61 ---@field highlighter_query vim.treesitter.highlighter.Query 62 63 ---@nodoc 64 ---@class vim.treesitter.highlighter 65 ---@field active table<integer,vim.treesitter.highlighter> 66 ---@field bufnr integer 67 ---@field private orig_spelloptions string 68 --- A map from window ID to highlight states. 69 --- This state is kept during rendering across each line update. 70 ---@field private _highlight_states vim.treesitter.highlighter.State[] 71 ---@field private _queries table<string,vim.treesitter.highlighter.Query> 72 ---@field _conceal_line boolean? 73 ---@field _conceal_checked table<integer, boolean> 74 ---@field tree vim.treesitter.LanguageTree 75 ---@field private redraw_count integer 76 --- A map from window ID to whether we are currently parsing that window asynchronously 77 ---@field parsing boolean 78 local TSHighlighter = { 79 active = {}, 80 } 81 82 TSHighlighter.__index = TSHighlighter 83 84 ---@nodoc 85 --- 86 --- Creates a highlighter for `tree`. 87 --- 88 ---@param tree vim.treesitter.LanguageTree parser object to use for highlighting 89 ---@param opts (table|nil) Configuration of the highlighter: 90 --- - queries table overwrite queries used by the highlighter 91 ---@return vim.treesitter.highlighter Created highlighter object 92 function TSHighlighter.new(tree, opts) 93 local self = setmetatable({}, TSHighlighter) 94 95 if type(tree:source()) ~= 'number' then 96 error('TSHighlighter can not be used with a string parser source.') 97 end 98 99 opts = opts or {} ---@type { queries: table<string,string> } 100 self.tree = tree 101 tree:register_cbs({ 102 on_detach = function() 103 self:on_detach() 104 end, 105 }) 106 107 -- Enable conceal_lines if query exists for lang and has conceal_lines metadata. 108 local function set_conceal_lines(lang) 109 if not self._conceal_line and self:get_query(lang):query() then 110 self._conceal_line = self:get_query(lang):query().has_conceal_line 111 end 112 end 113 114 tree:register_cbs({ 115 on_bytes = function(buf) 116 -- Clear conceal_lines marks whenever the buffer text changes. Marks are added 117 -- back as either the _conceal_line or on_win callback comes across them. 118 local hl = TSHighlighter.active[buf] 119 if hl and next(hl._conceal_checked) then 120 api.nvim_buf_clear_namespace(buf, ns, 0, -1) 121 hl._conceal_checked = {} 122 end 123 end, 124 on_changedtree = function(...) 125 self:on_changedtree(...) 126 end, 127 on_child_removed = function(child) 128 child:for_each_tree(function(t) 129 self:on_changedtree(t:included_ranges(true)) 130 end) 131 end, 132 on_child_added = function(child) 133 child:for_each_tree(function(t) 134 set_conceal_lines(t:lang()) 135 end) 136 end, 137 }, true) 138 139 local source = tree:source() 140 assert(type(source) == 'number') 141 142 self.bufnr = source 143 self.redraw_count = 0 144 self._conceal_checked = {} 145 self._queries = {} 146 self._highlight_states = {} 147 self.parsing = false 148 149 -- Queries for a specific language can be overridden by a custom 150 -- string query... if one is not provided it will be looked up by file. 151 if opts.queries then 152 for lang, query_string in pairs(opts.queries) do 153 self._queries[lang] = TSHighlighterQuery.new(lang, query_string) 154 set_conceal_lines(lang) 155 end 156 end 157 set_conceal_lines(tree:lang()) 158 self.orig_spelloptions = vim.bo[self.bufnr].spelloptions 159 160 vim.bo[self.bufnr].syntax = '' 161 vim.b[self.bufnr].ts_highlight = true 162 163 TSHighlighter.active[self.bufnr] = self 164 165 -- Tricky: if syntax hasn't been enabled, we need to reload color scheme 166 -- but use synload.vim rather than syntax.vim to not enable 167 -- syntax FileType autocmds. Later on we should integrate with the 168 -- `:syntax` and `set syntax=...` machinery properly. 169 -- Still need to ensure that syntaxset augroup exists, so that calling :destroy() 170 -- immediately afterwards will not error. 171 if vim.g.syntax_on ~= 1 then 172 vim.cmd.runtime({ 'syntax/synload.vim', bang = true }) 173 api.nvim_create_augroup('syntaxset', { clear = false }) 174 end 175 176 vim._with({ buf = self.bufnr }, function() 177 vim.opt_local.spelloptions:append('noplainbuffer') 178 end) 179 180 return self 181 end 182 183 --- @nodoc 184 --- Removes all internal references to the highlighter 185 function TSHighlighter:destroy() 186 TSHighlighter.active[self.bufnr] = nil 187 188 if api.nvim_buf_is_loaded(self.bufnr) then 189 vim.bo[self.bufnr].spelloptions = self.orig_spelloptions 190 vim.b[self.bufnr].ts_highlight = nil 191 api.nvim_buf_clear_namespace(self.bufnr, ns, 0, -1) 192 if vim.g.syntax_on == 1 then 193 -- FileType autocmds commonly assume curbuf is the target buffer, so nvim_buf_call. 194 api.nvim_buf_call(self.bufnr, function() 195 api.nvim_exec_autocmds( 196 'FileType', 197 { group = 'syntaxset', buffer = self.bufnr, modeline = false } 198 ) 199 end) 200 end 201 end 202 end 203 204 ---@param srow integer 205 ---@param erow integer exclusive 206 ---@private 207 function TSHighlighter:prepare_highlight_states(srow, erow) 208 self._highlight_states = {} 209 210 self.tree:for_each_tree(function(tstree, tree) 211 if not tstree then 212 return 213 end 214 215 local root_node = tstree:root() 216 local root_start_row, _, root_end_row, _ = root_node:range() 217 218 -- Only consider trees within the visible range 219 if root_start_row > erow or root_end_row < srow then 220 return 221 end 222 223 local hl_query = self:get_query(tree:lang()) 224 -- Some injected languages may not have highlight queries. 225 if not hl_query:query() then 226 return 227 end 228 229 -- _highlight_states should be a list so that the highlights are added in the same order as 230 -- for_each_tree traversal. This ensures that parents' highlight don't override children's. 231 table.insert(self._highlight_states, { 232 tstree = tstree, 233 next_row = 0, 234 next_col = 0, 235 iter = nil, 236 highlighter_query = hl_query, 237 }) 238 end) 239 end 240 241 ---@param fn fun(state: vim.treesitter.highlighter.State) 242 ---@package 243 function TSHighlighter:for_each_highlight_state(fn) 244 for _, state in ipairs(self._highlight_states) do 245 fn(state) 246 end 247 end 248 249 ---@package 250 function TSHighlighter:on_detach() 251 self:destroy() 252 end 253 254 ---@package 255 ---@param changes Range6[] 256 function TSHighlighter:on_changedtree(changes) 257 for _, ch in ipairs(changes) do 258 api.nvim__redraw({ buf = self.bufnr, range = { ch[1], ch[4] + 1 }, flush = false }) 259 -- Only invalidate the _conceal_checked range if _conceal_line is set and 260 -- ch[4] is not UINT32_MAX (empty range on first changedtree). 261 if ch[4] == 2 ^ 32 - 1 then 262 self._conceal_checked = {} 263 end 264 for i = ch[1], self._conceal_line and ch[4] ~= 2 ^ 32 - 1 and ch[4] or 0 do 265 self._conceal_checked[i] = false 266 end 267 end 268 end 269 270 --- Gets the query used for @param lang 271 ---@nodoc 272 ---@param lang string Language used by the highlighter. 273 ---@return vim.treesitter.highlighter.Query 274 function TSHighlighter:get_query(lang) 275 if not self._queries[lang] then 276 local success, result = pcall(TSHighlighterQuery.new, lang) 277 if not success then 278 self:destroy() 279 error(result) 280 end 281 self._queries[lang] = result 282 end 283 284 return self._queries[lang] 285 end 286 287 --- @param match TSQueryMatch 288 --- @param bufnr integer 289 --- @param capture integer 290 --- @param metadata vim.treesitter.query.TSMetadata 291 --- @return string? 292 local function get_url(match, bufnr, capture, metadata) 293 ---@type string|number|nil 294 local url = metadata[capture] and metadata[capture].url 295 296 if not url or type(url) == 'string' then 297 return url 298 end 299 300 local captures = match:captures() 301 302 if not captures[url] then 303 return 304 end 305 306 -- Assume there is only one matching node. If there is more than one, take the URL 307 -- from the first. 308 local other_node = captures[url][1] 309 310 return vim.treesitter.get_node_text(other_node, bufnr, { 311 metadata = metadata[url], 312 }) 313 end 314 315 --- @param capture_name string 316 --- @return boolean?, integer 317 local function get_spell(capture_name) 318 if capture_name == 'spell' then 319 return true, 0 320 elseif capture_name == 'nospell' then 321 -- Give nospell a higher priority so it always overrides spell captures. 322 return false, 1 323 end 324 return nil, 0 325 end 326 327 ---@param self vim.treesitter.highlighter 328 ---@param buf integer 329 ---@param range_start_row integer 330 ---@param range_start_col integer 331 ---@param range_end_row integer 332 ---@param range_end_col integer 333 ---@param on_spell boolean 334 ---@param on_conceal boolean 335 local function on_range_impl( 336 self, 337 buf, 338 range_start_row, 339 range_start_col, 340 range_end_row, 341 range_end_col, 342 on_spell, 343 on_conceal 344 ) 345 if self._conceal_line then 346 range_start_col = 0 347 if range_end_col ~= 0 then 348 range_end_row = range_end_row + 1 349 range_end_col = 0 350 end 351 end 352 for i = range_start_row, range_end_row - 1 do 353 self._conceal_checked[i] = self._conceal_line or nil 354 end 355 356 local MAX_ROW = 2147483647 -- sentinel for skipping to the end of file 357 local skip_until_row = MAX_ROW 358 local skip_until_col = 0 359 360 local subtree_counter = 0 361 self:for_each_highlight_state(function(state) 362 subtree_counter = subtree_counter + 1 363 local root_node = state.tstree:root() 364 ---@type { [1]: integer, [2]: integer, [3]: integer, [4]: integer } 365 local root_range = { root_node:range() } 366 367 if 368 not Range.intercepts( 369 root_range, 370 { range_start_row, range_start_col, range_end_row, range_end_col } 371 ) 372 then 373 if cmp_lt(root_range[1], root_range[2], skip_until_row, skip_until_col) then 374 skip_until_row = root_range[1] 375 skip_until_col = root_range[2] 376 end 377 return 378 end 379 380 local tree_region = state.tstree:included_ranges(true) 381 382 local next_row = state.next_row 383 local next_col = state.next_col 384 385 if state.iter == nil or cmp_lt(next_row, next_col, range_start_row, range_start_col) then 386 -- Mainly used to skip over folds 387 388 -- TODO(lewis6991): Creating a new iterator loses the cached predicate results for query 389 -- matches. Move this logic inside iter_captures() so we can maintain the cache. 390 state.iter = state.highlighter_query:query():iter_captures( 391 root_node, 392 self.bufnr, 393 range_start_row, 394 root_range[3], 395 { start_col = range_start_col, end_col = root_range[4] } 396 ) 397 end 398 399 local captures = state.highlighter_query:query().captures 400 401 while cmp_lt(next_row, next_col, range_end_row, range_end_col) do 402 local capture, node, metadata, match = state.iter(range_end_row, range_end_col) 403 if not node then 404 next_row = math.huge 405 next_col = math.huge 406 break 407 end 408 409 local outer_range = vim.treesitter.get_range(node, buf, metadata and metadata[capture]) 410 if cmp_lt(next_row, next_col, outer_range[1], outer_range[2]) then 411 next_row = outer_range[1] 412 next_col = outer_range[2] 413 end 414 415 if not capture then 416 break 417 end 418 419 for _, range in ipairs(tree_region) do 420 local intersection = Range.intersection(range, outer_range) 421 if intersection then 422 local start_row, start_col, end_row, end_col = Range.unpack4(intersection) 423 424 local hl = state.highlighter_query:get_hl_from_capture(capture) 425 426 local capture_name = captures[capture] 427 428 local spell, spell_pri_offset = get_spell(capture_name) 429 430 -- The "priority" attribute can be set at the pattern level or on a particular capture 431 local priority = ( 432 tonumber(metadata.priority or metadata[capture] and metadata[capture].priority) 433 or vim.hl.priorities.treesitter 434 ) + spell_pri_offset 435 436 -- The "conceal" attribute can be set at the pattern level or on a particular capture 437 local conceal = metadata.conceal or metadata[capture] and metadata[capture].conceal 438 439 local url = get_url(match, buf, capture, metadata) 440 441 if hl and not on_conceal and (not on_spell or spell ~= nil) then 442 -- Workaround for #35814: ensure the range is within buffer bounds, 443 -- allowing the last line if end_col is 0. 444 -- TODO(skewb1k): investigate a proper concurrency-safe handling of extmarks. 445 if (end_row + (end_col > 0 and 1 or 0)) <= api.nvim_buf_line_count(buf) then 446 api.nvim_buf_set_extmark(buf, ns, start_row, start_col, { 447 end_row = end_row, 448 end_col = end_col, 449 hl_group = hl, 450 ephemeral = true, 451 priority = priority, 452 conceal = conceal, 453 spell = spell, 454 url = url, 455 _subpriority = subtree_counter, 456 }) 457 end 458 end 459 460 if 461 (metadata.conceal_lines or metadata[capture] and metadata[capture].conceal_lines) 462 and #api.nvim_buf_get_extmarks(buf, ns, { start_row, 0 }, { start_row, 0 }, {}) == 0 463 then 464 api.nvim_buf_set_extmark(buf, ns, start_row, 0, { 465 end_line = end_row, 466 conceal_lines = '', 467 }) 468 end 469 end 470 end 471 end 472 473 state.next_row = next_row 474 state.next_col = next_col 475 if cmp_lt(next_row, next_col, skip_until_row, skip_until_col) then 476 skip_until_row = next_row 477 skip_until_col = next_col 478 end 479 end) 480 return skip_until_row, skip_until_col 481 end 482 483 ---@private 484 ---@param buf integer 485 ---@param br integer 486 ---@param bc integer 487 ---@param er integer 488 ---@param ec integer 489 function TSHighlighter._on_range(_, _, buf, br, bc, er, ec, _) 490 local self = TSHighlighter.active[buf] 491 if not self then 492 return 493 end 494 495 return on_range_impl(self, buf, br, bc, er, ec, false, false) 496 end 497 498 ---@private 499 ---@param buf integer 500 ---@param srow integer 501 ---@param erow integer 502 function TSHighlighter._on_spell_nav(_, _, buf, srow, _, erow, _) 503 local self = TSHighlighter.active[buf] 504 if not self then 505 return 506 end 507 508 -- Do not affect potentially populated highlight state. Here we just want a temporary 509 -- empty state so the C code can detect whether the region should be spell checked. 510 local highlight_states = self._highlight_states 511 local search_erow = math.max(erow, srow + 1) 512 self:prepare_highlight_states(srow, erow) 513 514 on_range_impl(self, buf, srow, 0, search_erow, 0, true, false) 515 self._highlight_states = highlight_states 516 end 517 518 ---@private 519 ---@param buf integer 520 ---@param row integer 521 function TSHighlighter._on_conceal_line(_, _, buf, row) 522 local self = TSHighlighter.active[buf] 523 if not self or not self._conceal_line or self._conceal_checked[row] then 524 return 525 end 526 527 -- Do not affect potentially populated highlight state. 528 local highlight_states = self._highlight_states 529 self.tree:parse({ row, row }) 530 self:prepare_highlight_states(row, row) 531 on_range_impl(self, buf, row, 0, row + 1, 0, false, true) 532 self._highlight_states = highlight_states 533 end 534 535 ---@private 536 ---@param buf integer 537 ---@param topline integer 538 ---@param botline integer 539 function TSHighlighter._on_win(_, _, buf, topline, botline) 540 local self = TSHighlighter.active[buf] 541 if not self then 542 return false 543 end 544 if not self.parsing then 545 self.redraw_count = self.redraw_count + 1 546 self:prepare_highlight_states(topline, botline) 547 else 548 self:for_each_highlight_state(function(state) 549 state.iter = nil 550 state.next_row = 0 551 state.next_col = 0 552 end) 553 end 554 return next(self._highlight_states) ~= nil 555 end 556 557 function TSHighlighter._on_start() 558 local buf_ranges = {} ---@type table<integer, Range[]> 559 for _, win in ipairs(api.nvim_tabpage_list_wins(0)) do 560 local buf = api.nvim_win_get_buf(win) 561 if TSHighlighter.active[buf] then 562 if not buf_ranges[buf] then 563 buf_ranges[buf] = {} 564 end 565 local topline = vim.fn.line('w0', win) - 1 566 -- +1 because w$ is the last completely displayed line (w_botline - 1), which may be -1 of the 567 -- last line that is at least partially visible. 568 local botline = vim.fn.line('w$', win) + 1 569 table.insert(buf_ranges[buf], { topline, botline }) 570 end 571 end 572 for buf, ranges in pairs(buf_ranges) do 573 local highlighter = TSHighlighter.active[buf] 574 if not highlighter.parsing then 575 table.sort(ranges, function(a, b) 576 return a[1] < b[1] 577 end) 578 highlighter.parsing = highlighter.parsing 579 or nil 580 == highlighter.tree:parse(ranges, function(_, trees) 581 if trees and highlighter.parsing then 582 highlighter.parsing = false 583 api.nvim__redraw({ buf = buf, valid = false, flush = false }) 584 end 585 end) 586 end 587 end 588 end 589 590 api.nvim_set_decoration_provider(ns, { 591 on_win = TSHighlighter._on_win, 592 on_start = TSHighlighter._on_start, 593 on_range = TSHighlighter._on_range, 594 _on_spell_nav = TSHighlighter._on_spell_nav, 595 _on_conceal_line = TSHighlighter._on_conceal_line, 596 }) 597 598 return TSHighlighter