fix(decor): allow adding providers during redraw

Fixes: #26652

6 files changed, 83 insertions, 93 deletions

diff --git a/src/nvim/api/extmark.c b/src/nvim/api/extmark.c
index fa3c5afcc6..b2ddef1a43 100644
--- a/src/nvim/api/extmark.c
+++ b/src/nvim/api/extmark.c
@@ -1050,7 +1050,7 @@ void nvim_set_decoration_provider(Integer ns_id, Dict(set_decoration_provider) *
     *v = LUA_NOREF;
   }
 
-  p->active = true;
+  p->state = kDecorProviderActive;
   p->hl_valid++;
   p->hl_cached = false;
 }
diff --git a/src/nvim/decoration_defs.h b/src/nvim/decoration_defs.h
index 4550ae0a42..d75a7bc242 100644
--- a/src/nvim/decoration_defs.h
+++ b/src/nvim/decoration_defs.h
@@ -124,7 +124,14 @@ typedef struct {
 
 typedef struct {
   NS ns_id;
-  bool active;
+
+  enum {
+    kDecorProviderActive = 1,
+    kDecorProviderWinDisabled = 2,
+    kDecorProviderRedrawDisabled = 3,
+    kDecorProviderDisabled = 4,
+  } state;
+
   LuaRef redraw_start;
   LuaRef redraw_buf;
   LuaRef redraw_win;
@@ -137,5 +144,3 @@ typedef struct {
 
   uint8_t error_count;
 } DecorProvider;
-
-typedef kvec_withinit_t(DecorProvider *, 4) DecorProviders;
diff --git a/src/nvim/decoration_provider.c b/src/nvim/decoration_provider.c
index e160a563f3..6a7ecf102d 100644
--- a/src/nvim/decoration_provider.c
+++ b/src/nvim/decoration_provider.c
@@ -26,7 +26,7 @@ enum { DP_MAX_ERROR = 3, };
 static kvec_t(DecorProvider) decor_providers = KV_INITIAL_VALUE;
 
 #define DECORATION_PROVIDER_INIT(ns_id) (DecorProvider) \
-  { ns_id, false, LUA_NOREF, LUA_NOREF, \
+  { ns_id, kDecorProviderDisabled, LUA_NOREF, LUA_NOREF, \
     LUA_NOREF, LUA_NOREF, LUA_NOREF, \
     LUA_NOREF, -1, false, false, 0 }
 
@@ -37,7 +37,9 @@ static void decor_provider_error(DecorProvider *provider, const char *name, cons
   msg_schedule_semsg_multiline("Error in decoration provider %s.%s:\n%s", ns_name, name, msg);
 }
 
-static bool decor_provider_invoke(DecorProvider *provider, const char *name, LuaRef ref, Array args,
+// Note we pass in a provider index as this function may cause decor_providers providers to be
+// reallocated so we need to be careful with DecorProvider pointers
+static bool decor_provider_invoke(int provider_idx, const char *name, LuaRef ref, Array args,
                                   bool default_true)
 {
   Error err = ERROR_INIT;
@@ -48,6 +50,10 @@ static bool decor_provider_invoke(DecorProvider *provider, const char *name, Lua
   provider_active = false;
   textlock--;
 
+  // We get the provider here via an index in case the above call to nlua_call_ref causes
+  // decor_providers to be reallocated.
+  DecorProvider *provider = &kv_A(decor_providers, provider_idx);
+
   if (!ERROR_SET(&err)
       && api_object_to_bool(ret, "provider %s retval", default_true, &err)) {
     provider->error_count = 0;
@@ -59,7 +65,7 @@ static bool decor_provider_invoke(DecorProvider *provider, const char *name, Lua
     provider->error_count++;
 
     if (provider->error_count >= DP_MAX_ERROR) {
-      provider->active = false;
+      provider->state = kDecorProviderDisabled;
     }
   }
 
@@ -72,11 +78,7 @@ void decor_providers_invoke_spell(win_T *wp, int start_row, int start_col, int e
 {
   for (size_t i = 0; i < kv_size(decor_providers); i++) {
     DecorProvider *p = &kv_A(decor_providers, i);
-    if (!p->active) {
-      continue;
-    }
-
-    if (p->spell_nav != LUA_NOREF) {
+    if (p->state != kDecorProviderDisabled && p->spell_nav != LUA_NOREF) {
       MAXSIZE_TEMP_ARRAY(args, 6);
       ADD_C(args, INTEGER_OBJ(wp->handle));
       ADD_C(args, INTEGER_OBJ(wp->w_buffer->handle));
@@ -84,7 +86,7 @@ void decor_providers_invoke_spell(win_T *wp, int start_row, int start_col, int e
       ADD_C(args, INTEGER_OBJ(start_col));
       ADD_C(args, INTEGER_OBJ(end_row));
       ADD_C(args, INTEGER_OBJ(end_col));
-      decor_provider_invoke(p, "spell", p->spell_nav, args, true);
+      decor_provider_invoke((int)i, "spell", p->spell_nav, args, true);
     }
   }
 }
@@ -93,27 +95,15 @@ void decor_providers_invoke_spell(win_T *wp, int start_row, int start_col, int e
 ///
 /// @param[out] providers Decoration providers
 /// @param[out] err       Provider err
-void decor_providers_start(DecorProviders *providers)
+void decor_providers_start(void)
 {
-  kvi_init(*providers);
-
   for (size_t i = 0; i < kv_size(decor_providers); i++) {
     DecorProvider *p = &kv_A(decor_providers, i);
-    if (!p->active) {
-      continue;
-    }
-
-    bool active;
-    if (p->redraw_start != LUA_NOREF) {
+    if (p->state != kDecorProviderDisabled && p->redraw_start != LUA_NOREF) {
       MAXSIZE_TEMP_ARRAY(args, 2);
       ADD_C(args, INTEGER_OBJ((int)display_tick));
-      active = decor_provider_invoke(p, "start", p->redraw_start, args, true);
-    } else {
-      active = true;
-    }
-
-    if (active) {
-      kvi_push(*providers, p);
+      bool active = decor_provider_invoke((int)i, "start", p->redraw_start, args, true);
+      kv_A(decor_providers, i).state = active ? kDecorProviderActive : kDecorProviderRedrawDisabled;
     }
   }
 }
@@ -125,10 +115,8 @@ void decor_providers_start(DecorProviders *providers)
 /// @param      providers      Decoration providers
 /// @param[out] line_providers Enabled line providers to invoke in win_line
 /// @param[out] err            Provider error
-void decor_providers_invoke_win(win_T *wp, DecorProviders *providers,
-                                DecorProviders *line_providers)
+void decor_providers_invoke_win(win_T *wp)
 {
-  kvi_init(*line_providers);
   // this might change in the future
   // then we would need decor_state.running_decor_provider just like "on_line" below
   assert(kv_size(decor_state.active) == 0);
@@ -138,17 +126,21 @@ void decor_providers_invoke_win(win_T *wp, DecorProviders *providers,
                            ? wp->w_botline
                            : MAX(wp->w_topline + wp->w_height_inner, wp->w_botline)));
 
-  for (size_t k = 0; k < kv_size(*providers); k++) {
-    DecorProvider *p = kv_A(*providers, k);
-    if (p && p->redraw_win != LUA_NOREF) {
+  for (size_t i = 0; i < kv_size(decor_providers); i++) {
+    DecorProvider *p = &kv_A(decor_providers, i);
+    if (p->state == kDecorProviderWinDisabled) {
+      p->state = kDecorProviderActive;
+    }
+
+    if (p->state == kDecorProviderActive && p->redraw_win != LUA_NOREF) {
       MAXSIZE_TEMP_ARRAY(args, 4);
       ADD_C(args, WINDOW_OBJ(wp->handle));
       ADD_C(args, BUFFER_OBJ(wp->w_buffer->handle));
       // TODO(bfredl): we are not using this, but should be first drawn line?
       ADD_C(args, INTEGER_OBJ(wp->w_topline - 1));
       ADD_C(args, INTEGER_OBJ(knownmax - 1));
-      if (decor_provider_invoke(p, "win", p->redraw_win, args, true)) {
-        kvi_push(*line_providers, p);
+      if (!decor_provider_invoke((int)i, "win", p->redraw_win, args, true)) {
+        kv_A(decor_providers, i).state = kDecorProviderWinDisabled;
       }
     }
   }
@@ -161,21 +153,21 @@ void decor_providers_invoke_win(win_T *wp, DecorProviders *providers,
 /// @param      row       Row to invoke line callback for
 /// @param[out] has_decor Set when at least one provider invokes a line callback
 /// @param[out] err       Provider error
-void decor_providers_invoke_line(win_T *wp, DecorProviders *providers, int row, bool *has_decor)
+void decor_providers_invoke_line(win_T *wp, int row, bool *has_decor)
 {
   decor_state.running_decor_provider = true;
-  for (size_t k = 0; k < kv_size(*providers); k++) {
-    DecorProvider *p = kv_A(*providers, k);
-    if (p && p->redraw_line != LUA_NOREF) {
+  for (size_t i = 0; i < kv_size(decor_providers); i++) {
+    DecorProvider *p = &kv_A(decor_providers, i);
+    if (p->state == kDecorProviderActive && p->redraw_line != LUA_NOREF) {
       MAXSIZE_TEMP_ARRAY(args, 3);
       ADD_C(args, WINDOW_OBJ(wp->handle));
       ADD_C(args, BUFFER_OBJ(wp->w_buffer->handle));
       ADD_C(args, INTEGER_OBJ(row));
-      if (decor_provider_invoke(p, "line", p->redraw_line, args, true)) {
+      if (decor_provider_invoke((int)i, "line", p->redraw_line, args, true)) {
         *has_decor = true;
       } else {
         // return 'false' or error: skip rest of this window
-        kv_A(*providers, k) = NULL;
+        kv_A(decor_providers, i).state = kDecorProviderWinDisabled;
       }
 
       hl_check_ns();
@@ -189,15 +181,15 @@ void decor_providers_invoke_line(win_T *wp, DecorProviders *providers, int row,
 /// @param      buf       Buffer
 /// @param      providers Decoration providers
 /// @param[out] err       Provider error
-void decor_providers_invoke_buf(buf_T *buf, DecorProviders *providers)
+void decor_providers_invoke_buf(buf_T *buf)
 {
-  for (size_t i = 0; i < kv_size(*providers); i++) {
-    DecorProvider *p = kv_A(*providers, i);
-    if (p && p->redraw_buf != LUA_NOREF) {
+  for (size_t i = 0; i < kv_size(decor_providers); i++) {
+    DecorProvider *p = &kv_A(decor_providers, i);
+    if (p->state == kDecorProviderActive && p->redraw_buf != LUA_NOREF) {
       MAXSIZE_TEMP_ARRAY(args, 2);
       ADD_C(args, BUFFER_OBJ(buf->handle));
       ADD_C(args, INTEGER_OBJ((int64_t)display_tick));
-      decor_provider_invoke(p, "buf", p->redraw_buf, args, true);
+      decor_provider_invoke((int)i, "buf", p->redraw_buf, args, true);
     }
   }
 }
@@ -207,14 +199,15 @@ void decor_providers_invoke_buf(buf_T *buf, DecorProviders *providers)
 /// @param      providers   Decoration providers
 /// @param      displaytick Display tick
 /// @param[out] err         Provider error
-void decor_providers_invoke_end(DecorProviders *providers)
+void decor_providers_invoke_end(void)
 {
-  for (size_t i = 0; i < kv_size(*providers); i++) {
-    DecorProvider *p = kv_A(*providers, i);
-    if (p && p->active && p->redraw_end != LUA_NOREF) {
+  for (size_t i = 0; i < kv_size(decor_providers); i++) {
+    DecorProvider *p = &kv_A(decor_providers, i);
+    if (p->state != kDecorProviderDisabled && p->redraw_end != LUA_NOREF) {
       MAXSIZE_TEMP_ARRAY(args, 1);
       ADD_C(args, INTEGER_OBJ((int)display_tick));
-      decor_provider_invoke(p, "end", p->redraw_end, args, true);
+      decor_provider_invoke((int)i, "end", p->redraw_end, args, true);
+      kv_A(decor_providers, i).state = kDecorProviderActive;
     }
   }
   decor_check_to_be_deleted();
@@ -227,10 +220,8 @@ void decor_providers_invoke_end(DecorProviders *providers)
 /// like highlight_changed() (throttled to the next redraw or mode change)
 void decor_provider_invalidate_hl(void)
 {
-  size_t len = kv_size(decor_providers);
-  for (size_t i = 0; i < len; i++) {
-    DecorProvider *item = &kv_A(decor_providers, i);
-    item->hl_cached = false;
+  for (size_t i = 0; i < kv_size(decor_providers); i++) {
+    kv_A(decor_providers, i).hl_cached = false;
   }
 
   if (ns_hl_active) {
@@ -242,14 +233,11 @@ void decor_provider_invalidate_hl(void)
 DecorProvider *get_decor_provider(NS ns_id, bool force)
 {
   assert(ns_id > 0);
-  size_t i;
   size_t len = kv_size(decor_providers);
-  for (i = 0; i < len; i++) {
-    DecorProvider *item = &kv_A(decor_providers, i);
-    if (item->ns_id == ns_id) {
-      return item;
-    } else if (item->ns_id > ns_id) {
-      break;
+  for (size_t i = 0; i < len; i++) {
+    DecorProvider *p = &kv_A(decor_providers, i);
+    if (p->ns_id == ns_id) {
+      return p;
     }
   }
 
@@ -257,16 +245,7 @@ DecorProvider *get_decor_provider(NS ns_id, bool force)
     return NULL;
   }
 
-  // Adding a new provider, so allocate room in the vector
-  (void)kv_a(decor_providers, len);
-  if (i < len) {
-    // New ns_id needs to be inserted between existing providers to maintain
-    // ordering, so shift other providers with larger ns_id
-    memmove(&kv_A(decor_providers, i + 1),
-            &kv_A(decor_providers, i),
-            (len - i) * sizeof(kv_a(decor_providers, i)));
-  }
-  DecorProvider *item = &kv_a(decor_providers, i);
+  DecorProvider *item = &kv_a(decor_providers, len);
   *item = DECORATION_PROVIDER_INIT(ns_id);
 
   return item;
@@ -283,7 +262,7 @@ void decor_provider_clear(DecorProvider *p)
   NLUA_CLEAR_REF(p->redraw_line);
   NLUA_CLEAR_REF(p->redraw_end);
   NLUA_CLEAR_REF(p->spell_nav);
-  p->active = false;
+  p->state = kDecorProviderDisabled;
 }
 
 void decor_free_all_mem(void)
diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index 134e81b4b0..e34c22d907 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -945,7 +945,7 @@ static void win_line_start(win_T *wp, winlinevars_T *wlv, bool save_extra)
 ///
 /// @return             the number of last row the line occupies.
 int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, bool number_only, spellvars_T *spv,
-             foldinfo_T foldinfo, DecorProviders *providers)
+             foldinfo_T foldinfo)
 {
   winlinevars_T wlv;                  // variables passed between functions
 
@@ -1081,7 +1081,7 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, bool number_onl
 
     has_decor = decor_redraw_line(wp, lnum - 1, &decor_state);
 
-    decor_providers_invoke_line(wp, providers, lnum - 1, &has_decor);
+    decor_providers_invoke_line(wp, lnum - 1, &has_decor);
 
     if (has_decor) {
       extra_check = true;
diff --git a/src/nvim/drawscreen.c b/src/nvim/drawscreen.c
index a6e483ab54..21e8438367 100644
--- a/src/nvim/drawscreen.c
+++ b/src/nvim/drawscreen.c
@@ -552,8 +552,7 @@ int update_screen(void)
 
   ui_comp_set_screen_valid(true);
 
-  DecorProviders providers;
-  decor_providers_start(&providers);
+  decor_providers_start();
 
   // "start" callback could have changed highlights for global elements
   if (win_check_ns_hl(NULL)) {
@@ -600,7 +599,7 @@ int update_screen(void)
       }
 
       if (buf->b_mod_tick_decor < display_tick) {
-        decor_providers_invoke_buf(buf, &providers);
+        decor_providers_invoke_buf(buf);
         buf->b_mod_tick_decor = display_tick;
       }
     }
@@ -630,7 +629,7 @@ int update_screen(void)
         did_one = true;
         start_search_hl();
       }
-      win_update(wp, &providers);
+      win_update(wp);
     }
 
     // redraw status line and window bar after the window to minimize cursor movement
@@ -671,8 +670,7 @@ int update_screen(void)
   }
   did_intro = true;
 
-  decor_providers_invoke_end(&providers);
-  kvi_destroy(providers);
+  decor_providers_invoke_end();
 
   // either cmdline is cleared, not drawn or mode is last drawn
   cmdline_was_last_drawn = false;
@@ -1427,7 +1425,7 @@ static void draw_sep_connectors_win(win_T *wp)
 /// top: from first row to top_end (when scrolled down)
 /// mid: from mid_start to mid_end (update inversion or changed text)
 /// bot: from bot_start to last row (when scrolled up)
-static void win_update(win_T *wp, DecorProviders *providers)
+static void win_update(win_T *wp)
 {
   int top_end = 0;              // Below last row of the top area that needs
                                 // updating.  0 when no top area updating.
@@ -1494,8 +1492,7 @@ static void win_update(win_T *wp, DecorProviders *providers)
 
   decor_redraw_reset(wp, &decor_state);
 
-  DecorProviders line_providers;
-  decor_providers_invoke_win(wp, providers, &line_providers);
+  decor_providers_invoke_win(wp);
 
   if (win_redraw_signcols(wp)) {
     wp->w_lines_valid = 0;
@@ -2284,7 +2281,7 @@ static void win_update(win_T *wp, DecorProviders *providers)
         spellvars_T zero_spv = { 0 };
         row = win_line(wp, lnum, srow, wp->w_grid.rows, false,
                        foldinfo.fi_lines > 0 ? &zero_spv : &spv,
-                       foldinfo, &line_providers);
+                       foldinfo);
 
         if (foldinfo.fi_lines == 0) {
           wp->w_lines[idx].wl_folded = false;
@@ -2322,7 +2319,7 @@ static void win_update(win_T *wp, DecorProviders *providers)
         // text doesn't need to be drawn, but the number column does.
         foldinfo_T info = wp->w_p_cul && lnum == wp->w_cursor.lnum
                           ? cursorline_fi : fold_info(wp, lnum);
-        (void)win_line(wp, lnum, srow, wp->w_grid.rows, true, &spv, info, &line_providers);
+        (void)win_line(wp, lnum, srow, wp->w_grid.rows, true, &spv, info);
       }
 
       // This line does not need to be drawn, advance to the next one.
@@ -2344,7 +2341,7 @@ static void win_update(win_T *wp, DecorProviders *providers)
       wp->w_lines_valid = 0;
       wp->w_valid &= ~VALID_WCOL;
       decor_redraw_reset(wp, &decor_state);
-      decor_providers_invoke_win(wp, providers, &line_providers);
+      decor_providers_invoke_win(wp);
       continue;
     }
 
@@ -2419,7 +2416,7 @@ static void win_update(win_T *wp, DecorProviders *providers)
         spellvars_T zero_spv = { 0 };
         foldinfo_T zero_foldinfo = { 0 };
         row = win_line(wp, wp->w_botline, row, wp->w_grid.rows, false, &zero_spv,
-                       zero_foldinfo, &line_providers);
+                       zero_foldinfo);
       }
     } else if (dollar_vcol == -1) {
       wp->w_botline = lnum;
@@ -2443,8 +2440,6 @@ static void win_update(win_T *wp, DecorProviders *providers)
     set_empty_rows(wp, row);
   }
 
-  kvi_destroy(line_providers);
-
   if (wp->w_redr_type >= UPD_REDRAW_TOP) {
     draw_vsep_win(wp);
     draw_hsep_win(wp);
@@ -2485,7 +2480,7 @@ static void win_update(win_T *wp, DecorProviders *providers)
         // Don't update for changes in buffer again.
         int mod_set = curbuf->b_mod_set;
         curbuf->b_mod_set = false;
-        win_update(curwin, providers);
+        win_update(curwin);
         must_redraw = 0;
         curbuf->b_mod_set = mod_set;
       }
diff --git a/test/functional/ui/decorations_spec.lua b/test/functional/ui/decorations_spec.lua
index f1290bbdca..0472b72e26 100644
--- a/test/functional/ui/decorations_spec.lua
+++ b/test/functional/ui/decorations_spec.lua
@@ -687,6 +687,17 @@ describe('decorations providers', function()
       {18:Press ENTER or type command to continue}^ |
     ]]}
   end)
+
+   it('can add new providers during redraw #26652', function()
+    setup_provider [[
+    local ns = api.nvim_create_namespace('test_no_add')
+    function on_do(...)
+      api.nvim_set_decoration_provider(ns, {})
+    end
+    ]]
+
+    helpers.assert_alive()
+  end)
 end)
 
 local example_text = [[