Merge pull request #31844 from bfredl/iter_crash

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

5 files changed, 31 insertions, 7 deletions

diff --git a/src/nvim/decoration.c b/src/nvim/decoration.c
index 728a917c22..a9f3ba0c3b 100644
--- a/src/nvim/decoration.c
+++ b/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;
diff --git a/src/nvim/decoration.h b/src/nvim/decoration.h
index 1d268c982b..a2f4fefd45 100644
--- a/src/nvim/decoration.h
+++ b/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 });
diff --git a/src/nvim/decoration_provider.c b/src/nvim/decoration_provider.c
index e5d2658720..7c99fbf889 100644
--- a/src/nvim/decoration_provider.c
+++ b/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;
       }
diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index 3062b0f2a3..d5273ff3d1 100644
--- a/src/nvim/drawline.c
+++ b/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;
     }
diff --git a/src/nvim/extmark.c b/src/nvim/extmark.c
index 6119d838f9..79eea718f4 100644
--- a/src/nvim/extmark.c
+++ b/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;
 }