From eb37241d38ad35b9e6bfac6379dd10e60aa0350c Mon Sep 17 00:00:00 2001
From: Riley Bruins <ribru17@hotmail.com>
Date: Sun, 26 May 2024 10:27:12 -0700
Subject: [PATCH] fix(tohtml): properly handle multiple hl groups #29012

Problem: :TOhtml doesn't properly handle virtual text when it has
multiple highlight groups. It also improperly calculates position offset
for multi-byte virt_text characters.

Solution: Apply the `vim.api.nvim_strwidth` broadly to properly
calculate character offset, and handle the cases where the `hl` argument
can be a table of multiple hl groups.
---
 runtime/lua/tohtml.lua                 | 14 +++++++++-----
 test/functional/plugin/tohtml_spec.lua |  8 +++++++-
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/runtime/lua/tohtml.lua b/runtime/lua/tohtml.lua
index 0fc349e86d..5e145950b7 100644
--- a/runtime/lua/tohtml.lua
+++ b/runtime/lua/tohtml.lua
@@ -188,6 +188,8 @@ local background_color_cache = nil
 --- @type string?
 local foreground_color_cache = nil
 
+local len = vim.api.nvim_strwidth
+
 --- @see https://invisible-island.net/xterm/ctlseqs/ctlseqs.html#h3-Operating-System-Commands
 --- @param color "background"|"foreground"|integer
 --- @return string?
@@ -312,9 +314,12 @@ local function style_line_insert_virt_text(style_line, col, val)
 end
 
 --- @param state vim.tohtml.state
---- @param hl string|integer|nil
+--- @param hl string|integer|string[]|integer[]?
 --- @return nil|integer
 local function register_hl(state, hl)
+  if type(hl) == 'table' then
+    hl = hl[#hl]
+  end
   if type(hl) == 'nil' then
     return
   elseif type(hl) == 'string' then
@@ -537,7 +542,7 @@ local function _styletable_extmarks_virt_text(state, extmark, namespaces)
       else
         style_line_insert_virt_text(styletable[row + 1], col + 1, { i[1], hlid })
       end
-      virt_text_len = virt_text_len + #i[1]
+      virt_text_len = virt_text_len + len(i[1])
     end
     if extmark[4].virt_text_pos == 'overlay' then
       styletable_insert_range(state, row + 1, col + 1, row + 1, col + virt_text_len + 1, HIDE_ID)
@@ -782,7 +787,7 @@ local function styletable_statuscolumn(state)
       statuscolumn,
       { winid = state.winid, use_statuscol_lnum = row, highlights = true }
     )
-    local width = vim.api.nvim_strwidth(status.str)
+    local width = len(status.str)
     if width > minwidth then
       minwidth = width
     end
@@ -797,7 +802,7 @@ local function styletable_statuscolumn(state)
     for k, v in ipairs(hls) do
       local text = str:sub(v.start + 1, hls[k + 1] and hls[k + 1].start or nil)
       if k == #hls then
-        text = text .. (' '):rep(minwidth - vim.api.nvim_strwidth(str))
+        text = text .. (' '):rep(minwidth - len(str))
       end
       if text ~= '' then
         local hlid = register_hl(state, v.group)
@@ -817,7 +822,6 @@ local function styletable_listchars(state)
   local function utf8_sub(str, i, j)
     return vim.fn.strcharpart(str, i - 1, j and j - i + 1 or nil)
   end
-  local len = vim.api.nvim_strwidth
   --- @type table<string,string>
   local listchars = vim.opt_local.listchars:get()
   local ids = setmetatable({}, {
diff --git a/test/functional/plugin/tohtml_spec.lua b/test/functional/plugin/tohtml_spec.lua
index 200a5f34b2..be5bada901 100644
--- a/test/functional/plugin/tohtml_spec.lua
+++ b/test/functional/plugin/tohtml_spec.lua
@@ -287,7 +287,13 @@ describe(':TOhtml', function()
         0,
         { virt_text = { { 'foo' } }, virt_text_pos = 'overlay' }
       )
-      api.nvim_buf_set_extmark(0, ns, 2, 0, { virt_text = { { 'foo' } }, virt_text_pos = 'inline' })
+      api.nvim_buf_set_extmark(
+        0,
+        ns,
+        2,
+        0,
+        { virt_text = { { 'fo┊o', { 'Conceal', 'Comment' } } }, virt_text_pos = 'inline' }
+      )
       --api.nvim_buf_set_extmark(0,ns,3,0,{virt_text={{'foo'}},virt_text_pos='right_align'})
       run_tohtml_and_assert(screen)
     end)