From 5c6ee251a6f9557f4941de10347d5aa9b408e33b Mon Sep 17 00:00:00 2001
From: Riley Bruins <ribru17@hotmail.com>
Date: Mon, 5 May 2025 19:05:16 -0700
Subject: [PATCH] fix(treesitter): eliminate flicker for single windows #33842

This commit will eliminate flicker while editing one open window. It
works by querying previously calculated trees for highlight marks, in
the case that we are still asynchronously parsing in `on_win`.

It will not fully solve the case of flicker when there are multiple open
windows, since the parser will drop previously parsed injection trees
whenever it receives a call to parse a different region of the buffer.
This will require a refactor of `languagetree.lua`.

(cherry picked from commit 8d75910ef9689b601087cf9bc59ec2622771490a)
---
 runtime/lua/vim/treesitter/highlighter.lua | 39 ++++++++++++++++------
 1 file changed, 29 insertions(+), 10 deletions(-)

diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index 5e57124d63..fa8eea9467 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -131,6 +131,7 @@ function TSHighlighter.new(tree, opts)
   self.redraw_count = 0
   self._conceal_checked = {}
   self._queries = {}
+  self._highlight_states = {}
 
   -- Queries for a specific language can be overridden by a custom
   -- string query... if one is not provided it will be looked up by file.
@@ -460,19 +461,37 @@ end
 ---@param buf integer
 ---@param topline integer
 ---@param botline integer
-function TSHighlighter._on_win(_, _, buf, topline, botline)
+function TSHighlighter._on_win(_, win, buf, topline, botline)
   local self = TSHighlighter.active[buf]
-  if not self or self.parsing then
+  if not self then
     return false
   end
-  self.parsing = self.tree:parse({ topline, botline + 1 }, function(_, trees)
-    if trees and self.parsing then
-      self.parsing = false
-      api.nvim__redraw({ buf = buf, valid = false, flush = false })
-    end
-  end) == nil
-  self.redraw_count = self.redraw_count + 1
-  self:prepare_highlight_states(topline, botline)
+  self.parsing = self.parsing
+    or nil
+      == self.tree:parse({ topline, botline + 1 }, function(_, trees)
+        if trees and self.parsing then
+          self.parsing = false
+          api.nvim__redraw({ win = win, valid = false, flush = false })
+        end
+      end)
+  if not self.parsing then
+    self.redraw_count = self.redraw_count + 1
+    self:prepare_highlight_states(topline, botline)
+  else
+    self:for_each_highlight_state(function(state)
+      -- TODO(ribru17): Inefficient. Eventually all marks should be applied in on_buf, and all
+      -- non-folded ranges of each open window should be merged, and iterators should only be
+      -- created over those regions. This would also fix #31777.
+      --
+      -- Currently this is not possible because the parser discards previously parsed injection
+      -- trees upon parsing a different region.
+      --
+      -- It would also be nice if rather than re-querying extmarks for old trees, we could tell the
+      -- decoration provider to not clear previous ephemeral marks for this redraw cycle.
+      state.iter = nil
+      state.next_row = 0
+    end)
+  end
   return #self._highlight_states > 0
 end