Merge pull request #31844 from bfredl/iter_crash

fix(decoration): fix crash when on_lines decor provider modifies marktree

src/nvim/decoration.c

@@ -303,12 +303,24 @@ static void decor_free_inner(DecorVirtText *vt, uint32_t first_idx)
   }
 }
 
+/// Check if we are in a callback while drawing, which might invalidate the marktree iterator.
+///
+/// This should be called whenever a structural modification has been done to a
+/// marktree in a public API function (i e any change which adds or deletes marks).
+void decor_state_invalidate(buf_T *buf)
+{
+  if (decor_state.win && decor_state.win->w_buffer == buf) {
+    decor_state.itr_valid = false;
+  }
+}
+
 void decor_check_to_be_deleted(void)
 {
   assert(!decor_state.running_decor_provider);
   decor_free_inner(to_free_virt, to_free_sh);
   to_free_virt = NULL;
   to_free_sh = DECOR_ID_INVALID;
+  decor_state.win = NULL;
 }
 
 void decor_state_free(DecorState *state)
@@ -447,6 +459,8 @@ bool decor_redraw_start(win_T *wp, int top_row, DecorState *state)
 {
   buf_T *buf = wp->w_buffer;
   state->top_row = top_row;
+  state->itr_valid = true;
+
   if (!marktree_itr_get_overlap(buf->b_marktree, top_row, 0, state->itr)) {
     return false;
   }
@@ -489,7 +503,11 @@ bool decor_redraw_line(win_T *wp, int row, DecorState *state)
 
   if (state->row == -1) {
     decor_redraw_start(wp, row, state);
+  } else if (!state->itr_valid) {
+    marktree_itr_get(wp->w_buffer->b_marktree, row, 0, state->itr);
+    state->itr_valid = true;
   }
+
   state->row = row;
   state->col_until = -1;
   state->eol_col = -1;

src/nvim/decoration.h

@@ -96,6 +96,7 @@ typedef struct {
   TriState spell;
 
   bool running_decor_provider;
+  bool itr_valid;
 } DecorState;
 
 EXTERN DecorState decor_state INIT( = { 0 });

src/nvim/decoration_provider.c

@@ -155,7 +155,7 @@ void decor_providers_invoke_win(win_T *wp)
 /// @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, int row, bool *has_decor)
+void decor_providers_invoke_line(win_T *wp, int row)
 {
   decor_state.running_decor_provider = true;
   for (size_t i = 0; i < kv_size(decor_providers); i++) {
@@ -165,9 +165,7 @@ void decor_providers_invoke_line(win_T *wp, int row, bool *has_decor)
       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((int)i, "line", p->redraw_line, args, true)) {
-        *has_decor = true;
-      } else {
+      if (!decor_provider_invoke((int)i, "line", p->redraw_line, args, true)) {
         // return 'false' or error: skip rest of this window
         kv_A(decor_providers, i).state = kDecorProviderWinDisabled;
       }

src/nvim/drawline.c

@@ -1051,12 +1051,12 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
       }
     }
 
-    has_decor = decor_redraw_line(wp, lnum - 1, &decor_state);
-
     if (!end_fill) {
-      decor_providers_invoke_line(wp, lnum - 1, &has_decor);
+      decor_providers_invoke_line(wp, lnum - 1);
     }
 
+    has_decor = decor_redraw_line(wp, lnum - 1, &decor_state);
+
     if (has_decor) {
       extra_check = true;
     }

src/nvim/extmark.c

@@ -95,6 +95,7 @@ void extmark_set(buf_T *buf, uint32_t ns_id, uint32_t *idp, int row, colnr_T col
   MTKey mark = { { row, col }, ns_id, id, flags, decor.data };
 
   marktree_put(buf->b_marktree, mark, end_row, end_col, end_right_gravity);
+  decor_state_invalidate(buf);
 
 revised:
   if (decor_flags || decor.ext) {
@@ -184,6 +185,8 @@ void extmark_del(buf_T *buf, MarkTreeIter *itr, MTKey key, bool restore)
     }
   }
 
+  decor_state_invalidate(buf);
+
   // TODO(bfredl): delete it from current undo header, opportunistically?
 }
 
@@ -237,6 +240,10 @@ bool extmark_clear(buf_T *buf, uint32_t ns_id, int l_row, colnr_T l_col, int u_r
     }
   }
 
+  if (marks_cleared_any) {
+    decor_state_invalidate(buf);
+  }
+
   return marks_cleared_any;
 }
 

test/functional/ui/decorations_spec.lua

@@ -744,6 +744,30 @@ describe('decorations providers', function()
     ]])
     eq(2, exec_lua([[return _G.cnt]]))
   end)
+
+  it('can do large changes to the marktree', function()
+    insert("line1 with a lot of text\nline2 with a lot of text")
+    setup_provider([[
+      function on_do(event, _, _, row)
+        if event == 'win' or (event == 'line' and row == 1) then
+          vim.api.nvim_buf_clear_namespace(0, ns1, 0, -1)
+          for i = 0,1 do
+            for j = 0,23 do
+              vim.api.nvim_buf_set_extmark(0, ns1, i, j, {hl_group='ErrorMsg', end_col = j+1})
+            end
+          end
+        end
+      end
+    ]])
+
+    -- Doesn't crash when modifying the marktree between line1 and line2
+    screen:expect([[
+      {2:line1 with a lot of text}                |
+      {2:line2 with a lot of tex^t}                |
+      {1:~                                       }|*5
+                                              |
+    ]])
+  end)
 end)
 
 local example_text = [[