feat(treesitter): vertical conceal support for highlighter

TSHighlighter now places marks for conceal_lines metadata. A new
internal decor provider callback _on_conceal_line was added that
instructs the highlighter to place conceal_lines marks whenever the
editor needs to know whether a line is concealed. The bundled markdown
queries use conceal_lines metadata to conceal code block fence lines.

3 files changed, 99 insertions, 39 deletions

diff --git a/runtime/lua/vim/_meta/api_keysets.lua b/runtime/lua/vim/_meta/api_keysets.lua
index e574c4c1b3..180032e7e9 100644
--- a/runtime/lua/vim/_meta/api_keysets.lua
+++ b/runtime/lua/vim/_meta/api_keysets.lua
@@ -236,6 +236,7 @@ error('Cannot require a meta file')
 --- @field on_end? fun(_: "end", tick: integer)
 --- @field _on_hl_def? fun(_: "hl_def")
 --- @field _on_spell_nav? fun(_: "spell_nav")
+--- @field _on_conceal_line? fun(_: "conceal_line")
 
 --- @class vim.api.keyset.set_extmark
 --- @field id? integer
diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index 475a1f0aa5..e0e2e2e091 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -67,6 +67,8 @@ end
 --- This state is kept during rendering across each line update.
 ---@field private _highlight_states vim.treesitter.highlighter.State[]
 ---@field private _queries table<string,vim.treesitter.highlighter.Query>
+---@field  _conceal_line boolean?
+---@field  _conceal_checked table<integer, boolean>
 ---@field tree vim.treesitter.LanguageTree
 ---@field private redraw_count integer
 ---@field parsing boolean true if we are parsing asynchronously
@@ -108,6 +110,11 @@ function TSHighlighter.new(tree, opts)
         self:on_changedtree(t:included_ranges(true))
       end)
     end,
+    on_child_added = function(child)
+      child:for_each_tree(function(t)
+        self._conceal_line = self._conceal_line or self:get_query(t:lang()):query().has_conceal_line
+      end)
+    end,
   }, true)
 
   local source = tree:source()
@@ -115,7 +122,7 @@ function TSHighlighter.new(tree, opts)
 
   self.bufnr = source
   self.redraw_count = 0
-  self._highlight_states = {}
+  self._conceal_checked = {}
   self._queries = {}
 
   -- Queries for a specific language can be overridden by a custom
@@ -123,8 +130,10 @@ function TSHighlighter.new(tree, opts)
   if opts.queries then
     for lang, query_string in pairs(opts.queries) do
       self._queries[lang] = TSHighlighterQuery.new(lang, query_string)
+      self._conceal_line = self._conceal_line or self._queries[lang]:query().has_conceal_line
     end
   end
+  self._conceal_line = self._conceal_line or self:get_query(tree:lang()):query().has_conceal_line
 
   self.orig_spelloptions = vim.bo[self.bufnr].spelloptions
 
@@ -141,7 +150,7 @@ function TSHighlighter.new(tree, opts)
   -- immediately afterwards will not error.
   if vim.g.syntax_on ~= 1 then
     vim.cmd.runtime({ 'syntax/synload.vim', bang = true })
-    vim.api.nvim_create_augroup('syntaxset', { clear = false })
+    api.nvim_create_augroup('syntaxset', { clear = false })
   end
 
   vim._with({ buf = self.bufnr }, function()
@@ -159,6 +168,7 @@ function TSHighlighter:destroy()
   if api.nvim_buf_is_loaded(self.bufnr) then
     vim.bo[self.bufnr].spelloptions = self.orig_spelloptions
     vim.b[self.bufnr].ts_highlight = nil
+    api.nvim_buf_clear_namespace(self.bufnr, ns, 0, -1)
     if vim.g.syntax_on == 1 then
       api.nvim_exec_autocmds(
         'FileType',
@@ -187,10 +197,9 @@ function TSHighlighter:prepare_highlight_states(srow, erow)
       return
     end
 
-    local highlighter_query = self:get_query(tree:lang())
-
+    local hl_query = self:get_query(tree:lang())
     -- Some injected languages may not have highlight queries.
-    if not highlighter_query:query() then
+    if not hl_query:query() then
       return
     end
 
@@ -200,7 +209,7 @@ function TSHighlighter:prepare_highlight_states(srow, erow)
       tstree = tstree,
       next_row = 0,
       iter = nil,
-      highlighter_query = highlighter_query,
+      highlighter_query = hl_query,
     })
   end)
 end
@@ -222,7 +231,15 @@ end
 ---@param changes Range6[]
 function TSHighlighter:on_changedtree(changes)
   for _, ch in ipairs(changes) do
-    api.nvim__redraw({ buf = self.bufnr, range = { ch[1], ch[4] + 1 }, flush = false })
+    api.nvim__redraw({ buf = self.bufnr, range = { ch[1], ch[4] }, flush = false })
+    -- Only invalidate the _conceal_checked range if _conceal_line is set and
+    -- ch[4] is not UINT32_MAX (empty range on first changedtree).
+    if ch[4] == 2 ^ 32 - 1 then
+      self._conceal_checked = {}
+    end
+    for i = ch[1], self._conceal_line and ch[4] ~= 2 ^ 32 - 1 and ch[4] or 0 do
+      self._conceal_checked[i] = false
+    end
   end
 end
 
@@ -286,8 +303,10 @@ end
 ---@param self vim.treesitter.highlighter
 ---@param buf integer
 ---@param line integer
----@param is_spell_nav boolean
-local function on_line_impl(self, buf, line, is_spell_nav)
+---@param on_spell boolean
+---@param on_conceal boolean
+local function on_line_impl(self, buf, line, on_spell, on_conceal)
+  self._conceal_checked[line] = true
   self:for_each_highlight_state(function(state)
     local root_node = state.tstree:root()
     local root_start_row, _, root_end_row, _ = root_node:range()
@@ -335,7 +354,7 @@ local function on_line_impl(self, buf, line, is_spell_nav)
 
         local url = get_url(match, buf, capture, metadata)
 
-        if hl and end_row >= line and (not is_spell_nav or spell ~= nil) then
+        if hl and end_row >= line and not on_conceal and (not on_spell or spell ~= nil) then
           api.nvim_buf_set_extmark(buf, ns, start_row, start_col, {
             end_line = end_row,
             end_col = end_col,
@@ -347,6 +366,16 @@ local function on_line_impl(self, buf, line, is_spell_nav)
             url = url,
           })
         end
+
+        if
+          (metadata.conceal_lines or metadata[capture] and metadata[capture].conceal_lines)
+          and #api.nvim_buf_get_extmarks(buf, ns, { start_row, 0 }, { start_row, 0 }, {}) == 0
+        then
+          api.nvim_buf_set_extmark(buf, ns, start_row, 0, {
+            end_line = end_row,
+            conceal_lines = '',
+          })
+        end
       end
 
       if start_row > line then
@@ -366,7 +395,7 @@ function TSHighlighter._on_line(_, _win, buf, line, _)
     return
   end
 
-  on_line_impl(self, buf, line, false)
+  on_line_impl(self, buf, line, false, false)
 end
 
 ---@private
@@ -385,13 +414,44 @@ function TSHighlighter._on_spell_nav(_, _, buf, srow, _, erow, _)
   self:prepare_highlight_states(srow, erow)
 
   for row = srow, erow do
-    on_line_impl(self, buf, row, true)
+    on_line_impl(self, buf, row, true, false)
   end
   self._highlight_states = highlight_states
 end
 
 ---@private
 ---@param buf integer
+---@param row integer
+function TSHighlighter._on_conceal_line(_, _, buf, row)
+  local self = TSHighlighter.active[buf]
+  if not self or not self._conceal_line or self._conceal_checked[row] then
+    return self and self._conceal_line
+  end
+
+  -- Do not affect potentially populated highlight state.
+  local highlight_states = self._highlight_states
+  self.tree:parse({ row, row })
+  self:prepare_highlight_states(row, row)
+  on_line_impl(self, buf, row, false, true)
+  self._highlight_states = highlight_states
+  return true
+end
+
+---@private
+--- Clear conceal_lines marks whenever we redraw for a buffer change. Marks are
+--- added back as either the _conceal_line or on_win callback comes across them.
+function TSHighlighter._on_buf(_, buf)
+  local self = TSHighlighter.active[buf]
+  if not self or not self._conceal_line then
+    return
+  end
+
+  api.nvim_buf_clear_namespace(buf, ns, 0, -1)
+  self._conceal_checked = {}
+end
+
+---@private
+---@param buf integer
 ---@param topline integer
 ---@param botline integer
 function TSHighlighter._on_win(_, _, buf, topline, botline)
@@ -413,7 +473,9 @@ end
 api.nvim_set_decoration_provider(ns, {
   on_win = TSHighlighter._on_win,
   on_line = TSHighlighter._on_line,
+  on_buf = TSHighlighter._on_buf,
   _on_spell_nav = TSHighlighter._on_spell_nav,
+  _on_conceal_line = TSHighlighter._on_conceal_line,
 })
 
 return TSHighlighter
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index d26aa8e604..f43c087457 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -10,6 +10,20 @@ local EXTENDS_FORMAT = '^;+%s*extends%s*$'
 
 local M = {}
 
+---@nodoc
+---Parsed query, see |vim.treesitter.query.parse()|
+---
+---@class vim.treesitter.Query
+---@field lang string parser language name
+---@field captures string[] list of (unique) capture names defined in query
+---@field info vim.treesitter.QueryInfo query context (e.g. captures, predicates, directives)
+---@field has_conceal_line boolean whether the query sets conceal_lines metadata
+---@field has_combined_injections boolean whether the query contains combined injections
+---@field query TSQuery userdata query object
+---@field private _processed_patterns table<integer, vim.treesitter.query.ProcessedPattern>
+local Query = {}
+Query.__index = Query
+
 local function is_directive(name)
   return string.sub(name, -1) == '!'
 end
@@ -28,15 +42,10 @@ end
 ---@field directives vim.treesitter.query.ProcessedDirective[]
 
 --- 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
+function Query:_process_patterns()
+  self._processed_patterns = {}
+
+  for k, pattern_list in pairs(self.info.patterns) do
     ---@type vim.treesitter.query.ProcessedPredicate[]
     local predicates = {}
     ---@type vim.treesitter.query.ProcessedDirective[]
@@ -50,7 +59,10 @@ 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
+          self._has_combined_injections = true
+        end
+        if vim.deep_equal(pattern, { 'set!', 'conceal_lines', '' }) then
+          self.has_conceal_line = true
         end
       else
         local should_match = true
@@ -62,25 +74,10 @@ local function process_patterns(patterns)
       end
     end
 
-    processed_patterns[k] = { predicates = predicates, directives = directives }
+    self._processed_patterns[k] = { predicates = predicates, directives = directives }
   end
-
-  return processed_patterns, has_combined
 end
 
----@nodoc
----Parsed query, see |vim.treesitter.query.parse()|
----
----@class vim.treesitter.Query
----@field lang string parser language name
----@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
-
 ---@package
 ---@see vim.treesitter.query.parse
 ---@param lang string
@@ -96,7 +93,7 @@ function Query.new(lang, ts_query)
     patterns = query_info.patterns,
   }
   self.captures = self.info.captures
-  self._processed_patterns, self.has_combined_injections = process_patterns(self.info.patterns)
+  self:_process_patterns()
   return self
 end