feat(extmarks): add virt_text_repeat_linebreak flag (#26625)

Problem:  Unable to predict which byte-offset to place virtual text to
          make it repeat visually in the wrapped part of a line.
Solution: Add a flag to nvim_buf_set_extmark() that causes virtual
          text to repeat in wrapped lines.

12 files changed, 55 insertions, 12 deletions

diff --git a/runtime/doc/api.txt b/runtime/doc/api.txt
index 7cb6051020..6a69bf7b41 100644
--- a/runtime/doc/api.txt
+++ b/runtime/doc/api.txt
@@ -2696,6 +2696,8 @@ nvim_buf_set_extmark({buffer}, {ns_id}, {line}, {col}, {*opts})
                     text is selected or hidden because of scrolling with
                     'nowrap' or 'smoothscroll'. Currently only affects
                     "overlay" virt_text.
+                  • virt_text_repeat_linebreak : repeat the virtual text on
+                    wrapped lines.
                   • hl_mode : control how highlights are combined with the
                     highlights of the text. Currently only affects virt_text
                     highlights, but might affect `hl_group` in later versions.
diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 454c9becc5..bc3902bc4c 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -339,11 +339,10 @@ The following changes to existing APIs or features add new behavior.
   In addition, |nvim_buf_get_extmarks()| has gained an "overlap" option to
   return such ranges even if they started before the specified position.
 
-• Extmarks can opt-out of precise undo tracking using the new "undo_restore"
-  flag to |nvim_buf_set_extmark()|
-
-• Extmarks can be automatically hidden or removed using the new "invalidate"
-  flag to |nvim_buf_set_extmark()|
+• The following flags were added to |nvim_buf_set_extmark()|:
+  - "undo_restore": opt-out extmarks of precise undo tracking.
+  - "invalidate": automatically hide or delete extmarks.
+  - "virt_text_repeat_linebreak": repeat virtual text on wrapped lines.
 
 • LSP hover and signature help now use Treesitter for highlighting of Markdown
   content.
diff --git a/runtime/lua/vim/_meta/api.lua b/runtime/lua/vim/_meta/api.lua
index c8c3e1c3af..ba18a16488 100644
--- a/runtime/lua/vim/_meta/api.lua
+++ b/runtime/lua/vim/_meta/api.lua
@@ -519,6 +519,8 @@ function vim.api.nvim_buf_line_count(buffer) end
 ---                 text is selected or hidden because of scrolling with
 ---                 'nowrap' or 'smoothscroll'. Currently only affects
 ---                 "overlay" virt_text.
+---               • virt_text_repeat_linebreak : repeat the virtual text on
+---                 wrapped lines.
 ---               • hl_mode : control how highlights are combined with the
 ---                 highlights of the text. Currently only affects virt_text
 ---                 highlights, but might affect `hl_group` in later versions.
diff --git a/runtime/lua/vim/_meta/api_keysets.lua b/runtime/lua/vim/_meta/api_keysets.lua
index 4ec8b03d30..bb1031b2fa 100644
--- a/runtime/lua/vim/_meta/api_keysets.lua
+++ b/runtime/lua/vim/_meta/api_keysets.lua
@@ -251,6 +251,7 @@ error('Cannot require a meta file')
 --- @field virt_text_pos? string
 --- @field virt_text_win_col? integer
 --- @field virt_text_hide? boolean
+--- @field virt_text_repeat_linebreak? boolean
 --- @field hl_eol? boolean
 --- @field hl_mode? string
 --- @field invalidate? boolean
diff --git a/src/nvim/api/extmark.c b/src/nvim/api/extmark.c
index b2ddef1a43..0cf16294b6 100644
--- a/src/nvim/api/extmark.c
+++ b/src/nvim/api/extmark.c
@@ -391,6 +391,8 @@ Array nvim_buf_get_extmarks(Buffer buffer, Integer ns_id, Object start, Object e
 ///                                  text is selected or hidden because of
 ///                                  scrolling with 'nowrap' or 'smoothscroll'.
 ///                                  Currently only affects "overlay" virt_text.
+///               - virt_text_repeat_linebreak : repeat the virtual text on
+///                                              wrapped lines.
 ///               - hl_mode : control how highlights are combined with the
 ///                           highlights of the text. Currently only affects
 ///                           virt_text highlights, but might affect `hl_group`
@@ -613,7 +615,8 @@ Integer nvim_buf_set_extmark(Buffer buffer, Integer ns_id, Integer line, Integer
   }
 
   hl.flags |= opts->hl_eol ? kSHHlEol : 0;
-  virt_text.flags |= opts->virt_text_hide ? kVTHide : 0;
+  virt_text.flags |= ((opts->virt_text_hide ? kVTHide : 0)
+                      | (opts->virt_text_repeat_linebreak ? kVTRepeatLinebreak : 0));
 
   if (HAS_KEY(opts, set_extmark, hl_mode)) {
     String str = opts->hl_mode;
diff --git a/src/nvim/api/keysets_defs.h b/src/nvim/api/keysets_defs.h
index fa9a2cb013..897d546f3d 100644
--- a/src/nvim/api/keysets_defs.h
+++ b/src/nvim/api/keysets_defs.h
@@ -33,6 +33,7 @@ typedef struct {
   String virt_text_pos;
   Integer virt_text_win_col;
   Boolean virt_text_hide;
+  Boolean virt_text_repeat_linebreak;
   Boolean hl_eol;
   String hl_mode;
   Boolean invalidate;
diff --git a/src/nvim/decoration.c b/src/nvim/decoration.c
index 8766f2ca76..02b95576ea 100644
--- a/src/nvim/decoration.c
+++ b/src/nvim/decoration.c
@@ -1061,11 +1061,11 @@ void decor_to_dict_legacy(Dictionary *dict, DecorInline decor, bool hl_name)
     Array chunks = virt_text_to_array(virt_text->data.virt_text, hl_name);
     PUT(*dict, "virt_text", ARRAY_OBJ(chunks));
     PUT(*dict, "virt_text_hide", BOOLEAN_OBJ(virt_text->flags & kVTHide));
+    PUT(*dict, "virt_text_repeat_linebreak", BOOLEAN_OBJ(virt_text->flags & kVTRepeatLinebreak));
     if (virt_text->pos == kVPosWinCol) {
       PUT(*dict, "virt_text_win_col", INTEGER_OBJ(virt_text->col));
     }
-    PUT(*dict, "virt_text_pos",
-        CSTR_TO_OBJ(virt_text_pos_str[virt_text->pos]));
+    PUT(*dict, "virt_text_pos", CSTR_TO_OBJ(virt_text_pos_str[virt_text->pos]));
     priority = virt_text->priority;
   }
 
diff --git a/src/nvim/decoration_defs.h b/src/nvim/decoration_defs.h
index d75a7bc242..93bc202b94 100644
--- a/src/nvim/decoration_defs.h
+++ b/src/nvim/decoration_defs.h
@@ -77,6 +77,7 @@ enum {
   kVTIsLines = 1,
   kVTHide = 2,
   kVTLinesAbove = 4,
+  kVTRepeatLinebreak = 8,
 };
 
 typedef struct DecorVirtText DecorVirtText;
diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index de607e912f..885d308514 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -286,10 +286,12 @@ static void draw_virt_text(win_T *wp, buf_T *buf, int col_off, int *end_col, int
       int vcol = item->draw_col - col_off;
       col = draw_virt_text_item(buf, item->draw_col, vt->data.virt_text,
                                 vt->hl_mode, max_col, vcol);
+      if (vt->pos == kVPosEndOfLine && do_eol) {
+        state->eol_col = col + 1;
+      }
     }
-    item->draw_col = INT_MIN;  // deactivate
-    if (vt && vt->pos == kVPosEndOfLine && do_eol) {
-      state->eol_col = col + 1;
+    if (!vt || !(vt->flags & kVTRepeatLinebreak)) {
+      item->draw_col = INT_MIN;  // deactivate
     }
 
     *end_col = MAX(*end_col, col);
@@ -1588,7 +1590,9 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, bool number_onl
           && wlv.vcol >= wp->w_virtcol)
          || number_only)
         && wlv.filler_todo <= 0) {
-      draw_virt_text(wp, buf, win_col_offset, &wlv.col, wlv.row);
+      if (!number_only) {
+        draw_virt_text(wp, buf, win_col_offset, &wlv.col, wlv.row);
+      }
       // don't clear anything after wlv.col
       win_put_linebuf(wp, wlv.row, 0, wlv.col, wlv.col, bg_attr, false);
       // Pretend we have finished updating the window.  Except when
diff --git a/test/functional/api/extmark_spec.lua b/test/functional/api/extmark_spec.lua
index 06eb938927..56383986f3 100644
--- a/test/functional/api/extmark_spec.lua
+++ b/test/functional/api/extmark_spec.lua
@@ -1549,6 +1549,7 @@ describe('API/extmarks', function()
       virt_lines_above = true,
       virt_lines_leftcol = true,
       virt_text = { { "text", "Macro" }, { "???" }, { "stack", { "Type", "Search" } } },
+      virt_text_repeat_linebreak = false,
       virt_text_hide = true,
       virt_text_pos = "right_align",
     } }, get_extmark_by_id(ns, marks[1], { details = true }))
@@ -1557,6 +1558,7 @@ describe('API/extmarks', function()
       right_gravity = true,
       priority = 0,
       virt_text = { { "", "Macro" }, { "", { "Type", "Search" } }, { "" } },
+      virt_text_repeat_linebreak = false,
       virt_text_hide = false,
       virt_text_pos = "win_col",
       virt_text_win_col = 1,
diff --git a/test/functional/ui/bufhl_spec.lua b/test/functional/ui/bufhl_spec.lua
index f7b8144ff9..e47ab273bd 100644
--- a/test/functional/ui/bufhl_spec.lua
+++ b/test/functional/ui/bufhl_spec.lua
@@ -644,6 +644,7 @@ describe('Buffer highlighting', function()
         virt_text = s1,
         -- other details
         right_gravity = true,
+        virt_text_repeat_linebreak = false,
         virt_text_pos = 'eol',
         virt_text_hide = false,
       }}}, get_extmarks(id1, {0,0}, {0, -1}, {details=true}))
@@ -656,6 +657,7 @@ describe('Buffer highlighting', function()
         virt_text = s2,
         -- other details
         right_gravity = true,
+        virt_text_repeat_linebreak = false,
         virt_text_pos = 'eol',
         virt_text_hide = false,
       }}}, get_extmarks(id1, {lastline,0}, {lastline, -1}, {details=true}))
diff --git a/test/functional/ui/decorations_spec.lua b/test/functional/ui/decorations_spec.lua
index 0472b72e26..be0f41ef6b 100644
--- a/test/functional/ui/decorations_spec.lua
+++ b/test/functional/ui/decorations_spec.lua
@@ -2094,6 +2094,32 @@ describe('extmark decorations', function()
                                                         |
     ]]}
   end)
+
+  it('virt_text_repeat_linebreak repeats virtual text on wrapped lines', function()
+    screen:try_resize(40, 5)
+    meths.set_option_value('breakindent', true, {})
+    insert(example_text)
+    meths.buf_set_extmark(0, ns, 1, 0, { virt_text = {{'│', 'NonText'}}, virt_text_pos = 'overlay', virt_text_repeat_linebreak = true })
+    meths.buf_set_extmark(0, ns, 1, 3, { virt_text = {{'│', 'NonText'}}, virt_text_pos = 'overlay', virt_text_repeat_linebreak = true })
+    command('norm gg')
+    screen:expect{grid=[[
+      ^for _,item in ipairs(items) do          |
+      {1:│}  {1:│}local text, hl_id_cell, count = unpa|
+      {1:│}  {1:│}ck(item)                            |
+          if hl_id_cell ~= nil then           |
+                                              |
+    ]]}
+    meths.buf_clear_namespace(0, ns, 0, -1)
+    meths.buf_set_extmark(0, ns, 1, 0, { virt_text = {{'│', 'NonText'}}, virt_text_repeat_linebreak = true, virt_text_win_col = 0 })
+    meths.buf_set_extmark(0, ns, 1, 0, { virt_text = {{'│', 'NonText'}}, virt_text_repeat_linebreak = true, virt_text_win_col = 2 })
+    screen:expect{grid=[[
+      ^for _,item in ipairs(items) do          |
+      {1:│} {1:│} local text, hl_id_cell, count = unpa|
+      {1:│} {1:│} ck(item)                            |
+          if hl_id_cell ~= nil then           |
+                                              |
+    ]]}
+  end)
 end)
 
 describe('decorations: inline virtual text', function()