From 2aa794826611cd4874badd2409928fd04a64a75f Mon Sep 17 00:00:00 2001
From: luukvbaal <luukvbaal@gmail.com>
Date: Wed, 7 May 2025 03:03:54 +0200
Subject: [PATCH] fix(treesitter): invalidate conceal_lines cache earlier
 #33875

Problem:  conceal_lines cache is invalidated in `on_buf`
          which is too late for code calculating text height after a
          buffer change but before a redraw (like `lsp/util.lua`).

Solution: Replace `on_buf` with `on_bytes` handler that invalidates
          the cache and clears the marks.
---
 runtime/lua/vim/treesitter/highlighter.lua    | 26 +++++++------------
 test/functional/treesitter/highlight_spec.lua | 24 +++++++++++++++++
 2 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index 7801bed65c..94dad7d4c1 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -109,6 +109,15 @@ function TSHighlighter.new(tree, opts)
   end
 
   tree:register_cbs({
+    on_bytes = function(buf)
+      -- Clear conceal_lines marks whenever the buffer text changes. Marks are added
+      -- back as either the _conceal_line or on_win callback comes across them.
+      local hl = TSHighlighter.active[buf]
+      if hl and next(hl._conceal_checked) then
+        api.nvim_buf_clear_namespace(buf, ns, 0, -1)
+        hl._conceal_checked = {}
+      end
+    end,
     on_changedtree = function(...)
       self:on_changedtree(...)
     end,
@@ -313,7 +322,7 @@ end
 ---@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._conceal_checked[line] = self._conceal_line and true or nil
   self:for_each_highlight_state(function(state)
     local root_node = state.tstree:root()
     local root_start_row, _, root_end_row, _ = root_node:range()
@@ -388,7 +397,6 @@ local function on_line_impl(self, buf, line, on_spell, on_conceal)
               api.nvim_buf_set_extmark(buf, ns, start_row, 0, {
                 end_line = end_row,
                 conceal_lines = '',
-                invalidate = true,
               })
             end
           end
@@ -453,19 +461,6 @@ function TSHighlighter._on_conceal_line(_, _, buf, row)
   self._highlight_states = highlight_states
 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
@@ -507,7 +502,6 @@ 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,
 })
diff --git a/test/functional/treesitter/highlight_spec.lua b/test/functional/treesitter/highlight_spec.lua
index f661fde503..06fab54388 100644
--- a/test/functional/treesitter/highlight_spec.lua
+++ b/test/functional/treesitter/highlight_spec.lua
@@ -1394,6 +1394,30 @@ printf('Hello World!');
       {8:126 }^                                    |
                                               |
     ]])
+    feed('ggdj')
+    command('set concealcursor=n')
+    screen:expect([[
+      {8:  2 }{25:^printf}{16:(}{26:'Hello World!'}{16:);}             |
+      {8:  4 }                                    |
+      {8:  6 }{25:printf}{16:(}{26:'Hello World!'}{16:);}             |
+      {8:  8 }                                    |
+      {8: 10 }{25:printf}{16:(}{26:'Hello World!'}{16:);}             |
+      {8: 12 }                                    |
+      {8: 14 }{25:printf}{16:(}{26:'Hello World!'}{16:);}             |
+      {8: 16 }                                    |
+      {8: 18 }{25:printf}{16:(}{26:'Hello World!'}{16:);}             |
+      {8: 20 }                                    |
+      {8: 22 }{25:printf}{16:(}{26:'Hello World!'}{16:);}             |
+      {8: 24 }                                    |
+      {8: 26 }{25:printf}{16:(}{26:'Hello World!'}{16:);}             |
+      {8: 28 }                                    |
+      {8: 30 }{25:printf}{16:(}{26:'Hello World!'}{16:);}             |
+                                              |
+    ]])
+    exec_lua(function()
+      vim.api.nvim_buf_set_lines(0, 0, -1, false, {})
+      assert(vim.api.nvim_win_text_height(0, {}).all == 1, 'line concealed')
+    end)
   end)
 end)