fix(treesitter): inconsistent highlight of multiline combined injection #32619

Problem:
Combined injections not entirely highlighted.

Solution:
Reapply layer highlights on each line.

runtime/lua/vim/treesitter/highlighter.lua

@@ -52,11 +52,14 @@ function TSHighlighterQuery:query()
   return self._query
 end
 
+---@alias MarkInfo { start_line: integer, start_col: integer, opts: vim.api.keyset.set_extmark }
+
 ---@class (private) vim.treesitter.highlighter.State
 ---@field tstree TSTree
 ---@field next_row integer
 ---@field iter vim.treesitter.highlighter.Iter?
 ---@field highlighter_query vim.treesitter.highlighter.Query
+---@field prev_marks MarkInfo[]
 
 ---@nodoc
 ---@class vim.treesitter.highlighter
@@ -229,6 +232,7 @@ function TSHighlighter:prepare_highlight_states(win, srow, erow)
       next_row = 0,
       iter = nil,
       highlighter_query = hl_query,
+      prev_marks = {},
     })
   end)
 end
@@ -320,6 +324,35 @@ local function get_spell(capture_name)
   return nil, 0
 end
 
+---Adds the mark to the buffer, clipped by the line.
+---Queues the remainder if the mark continues after the line.
+---@param m MarkInfo
+---@param buf integer
+---@param line integer
+---@param next_marks MarkInfo[]
+local function add_mark(m, buf, line, next_marks)
+  local cur_start_l = m.start_line
+  local cur_start_c = m.start_col
+  if cur_start_l < line then
+    cur_start_l = line
+    cur_start_c = 0
+  end
+
+  local cur_opts = m.opts
+  if cur_opts.end_line >= line + 1 then
+    cur_opts = vim.deepcopy(cur_opts, true)
+    cur_opts.end_line = line + 1
+    cur_opts.end_col = 0
+    table.insert(next_marks, m)
+  end
+
+  local empty = cur_opts.end_line < cur_start_l
+    or (cur_opts.end_line == cur_start_l and cur_opts.end_col <= cur_start_c)
+  if cur_start_l <= line and not empty then
+    api.nvim_buf_set_extmark(buf, ns, cur_start_l, cur_start_c, cur_opts)
+  end
+end
+
 ---@param self vim.treesitter.highlighter
 ---@param win integer
 ---@param buf integer
@@ -339,6 +372,12 @@ local function on_line_impl(self, win, buf, line, on_spell, on_conceal)
 
     local tree_region = state.tstree:included_ranges(true)
 
+    local next_marks = {}
+
+    for _, mark in ipairs(state.prev_marks) do
+      add_mark(mark, buf, line, next_marks)
+    end
+
     if state.iter == nil or state.next_row < line then
       -- Mainly used to skip over folds
 
@@ -383,7 +422,7 @@ local function on_line_impl(self, win, buf, line, on_spell, on_conceal)
             local url = get_url(match, buf, capture, metadata)
 
             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, {
+              local opts = {
                 end_line = end_row,
                 end_col = end_col,
                 hl_group = hl,
@@ -392,7 +431,9 @@ local function on_line_impl(self, win, buf, line, on_spell, on_conceal)
                 conceal = conceal,
                 spell = spell,
                 url = url,
-              })
+              }
+              local mark = { start_line = start_row, start_col = start_col, opts = opts }
+              add_mark(mark, buf, line, next_marks)
             end
 
             if
@@ -412,6 +453,8 @@ local function on_line_impl(self, win, buf, line, on_spell, on_conceal)
         state.next_row = outer_range_start_row
       end
     end
+
+    state.prev_marks = next_marks
   end)
 end
 

test/functional/treesitter/highlight_spec.lua

@@ -550,24 +550,13 @@ describe('treesitter highlighting (C)', function()
     screen:expect([=[
       {18:-- }{25:print}{16:(}{26:[[}                                                      |
       {18:--}{26: some}                                                          |
-      {18:-- random}                                                        |
-      {18:-- text}                                                          |
-      {18:-- here]])}                                                       |
+      {18:--}{26: random}                                                        |
+      {18:--}{26: text}                                                          |
+      {18:--}{26: here]]}{16:)}                                                       |
       ^                                                                 |
       {1:~                                                                }|*11
                                                                        |
     ]=])
-    -- NOTE: Once #31777 is fixed, this test case should be updated to the following:
-    -- screen:expect([=[
-    --   {18:-- }{25:print}{16:(}{26:[[}                                                      |
-    --   {18:--}{26: some}                                                          |
-    --   {18:--}{26: random}                                                        |
-    --   {18:--}{26: text}                                                          |
-    --   {18:--}{26: here]]}{16:)}                                                       |
-    --   ^                                                                   |
-    --   {1:~                                                                }|*11
-    --                                                                    |
-    -- ]=])
   end)
 
   it('supports complicated combined injections', function()