buffer_diff: Ensure that base text has finished parsing before completing update (#46054)

Follow-up to #46001 

That initial fix partly addressed the issue for diffs managed by the
`GitStore`, but not for other diffs (e.g. those managed by the
`ActionLog` or `CommitView`). The underlying issue is that we switched
to using a `Buffer` to represent the diff base text, and when updating
the diff we were calling `set_text` on this buffer and not waiting for
reparsing to finish. When the base text was represented by a series of
independent `BufferSnapshot`s, this wasn't an issue because we would
parse the base text in the background as part of computing the diff
update. This PR fixes the issue by waiting on reparsing to finish after
each call to `set_text`.

Release Notes:

- N/A

crates/acp_thread/src/diff.rs

@@ -212,12 +212,16 @@ impl PendingDiff {
                     )
                 })?
                 .await;
-            buffer_diff.update(cx, |diff, cx| {
-                diff.set_snapshot(update.clone(), &text_snapshot, cx);
-                diff.secondary_diff().unwrap().update(cx, |diff, cx| {
-                    diff.set_snapshot(update, &text_snapshot, cx);
-                });
+            let (task1, task2) = buffer_diff.update(cx, |diff, cx| {
+                let task1 = diff.set_snapshot(update.clone(), &text_snapshot, cx);
+                let task2 = diff
+                    .secondary_diff()
+                    .unwrap()
+                    .update(cx, |diff, cx| diff.set_snapshot(update, &text_snapshot, cx));
+                (task1, task2)
             })?;
+            task1.await;
+            task2.await;
             diff.update(cx, |diff, cx| {
                 if let Diff::Pending(diff) = diff {
                     diff.update_visible_ranges(cx);
@@ -380,17 +384,20 @@ async fn build_buffer_diff(
         })?
         .await;
 
-    secondary_diff.update(cx, |secondary_diff, cx| {
-        secondary_diff.language_changed(language.clone(), language_registry.clone(), cx);
-        secondary_diff.set_snapshot(update.clone(), &buffer, cx);
-    })?;
+    secondary_diff
+        .update(cx, |secondary_diff, cx| {
+            secondary_diff.language_changed(language.clone(), language_registry.clone(), cx);
+            secondary_diff.set_snapshot(update.clone(), &buffer, cx)
+        })?
+        .await;
 
     let diff = cx.new(|cx| BufferDiff::new(&buffer, cx))?;
     diff.update(cx, |diff, cx| {
         diff.language_changed(language, language_registry, cx);
-        diff.set_snapshot(update.clone(), &buffer, cx);
         diff.set_secondary_diff(secondary_diff);
-    })?;
+        diff.set_snapshot(update.clone(), &buffer, cx)
+    })?
+    .await;
     Ok(diff)
 }
 

crates/action_log/src/action_log.rs

@@ -401,10 +401,11 @@ impl ActionLog {
         if let Ok(update) = update {
             let update = update.await;
 
-            let diff_snapshot = diff.update(cx, |diff, cx| {
-                diff.set_snapshot(update.clone(), &buffer_snapshot, cx);
-                diff.snapshot(cx)
-            })?;
+            diff.update(cx, |diff, cx| {
+                diff.set_snapshot(update.clone(), &buffer_snapshot, cx)
+            })?
+            .await;
+            let diff_snapshot = diff.update(cx, |diff, cx| diff.snapshot(cx))?;
 
             unreviewed_edits = cx
                 .background_spawn({

crates/buffer_diff/src/buffer_diff.rs

@@ -1148,7 +1148,7 @@ impl BufferDiff {
             None,
             cx,
         ));
-        this.set_snapshot(inner, &buffer, cx);
+        this.set_snapshot(inner, &buffer, cx).detach();
         this
     }
 
@@ -1293,26 +1293,19 @@ impl BufferDiff {
         language_registry: Option<Arc<LanguageRegistry>>,
         cx: &mut Context<Self>,
     ) {
-        let base_text = self.inner.base_text.downgrade();
+        let fut = self.inner.base_text.update(cx, |base_text, cx| {
+            if let Some(language_registry) = language_registry {
+                base_text.set_language_registry(language_registry);
+            }
+            base_text.set_language(language, cx);
+            base_text.parsing_idle()
+        });
         cx.spawn(async move |this, cx| {
-            let fut = base_text
-                .update(cx, |base_text, cx| {
-                    if let Some(language_registry) = language_registry {
-                        base_text.set_language_registry(language_registry);
-                    }
-                    base_text.set_language(language, cx);
-                    base_text.parsing_idle()
-                })
-                .ok()?;
-
             fut.await;
-
             this.update(cx, |_, cx| {
                 cx.emit(BufferDiffEvent::LanguageChanged);
             })
-            .ok()?;
-
-            Some(())
+            .ok();
         })
         .detach();
     }
@@ -1322,7 +1315,7 @@ impl BufferDiff {
         new_state: BufferDiffUpdate,
         buffer: &text::BufferSnapshot,
         cx: &mut Context<Self>,
-    ) -> Option<Range<Anchor>> {
+    ) -> Task<Option<Range<Anchor>>> {
         self.set_snapshot_with_secondary(new_state, buffer, None, false, cx)
     }
 
@@ -1333,7 +1326,7 @@ impl BufferDiff {
         secondary_diff_change: Option<Range<Anchor>>,
         clear_pending_hunks: bool,
         cx: &mut Context<Self>,
-    ) -> Option<Range<Anchor>> {
+    ) -> Task<Option<Range<Anchor>>> {
         log::debug!("set snapshot with secondary {secondary_diff_change:?}");
 
         let old_snapshot = self.snapshot(cx);
@@ -1372,13 +1365,16 @@ impl BufferDiff {
 
         let state = &mut self.inner;
         state.base_text_exists = new_state.base_text_exists;
-        if update.base_text_changed {
+        let parsing_idle = if update.base_text_changed {
             state.base_text.update(cx, |base_text, cx| {
                 base_text.set_capability(Capability::ReadWrite, cx);
                 base_text.set_text(new_state.base_text.clone(), cx);
                 base_text.set_capability(Capability::ReadOnly, cx);
+                Some(base_text.parsing_idle())
             })
-        }
+        } else {
+            None
+        };
         state.hunks = new_state.hunks;
         if update.base_text_changed || clear_pending_hunks {
             if let Some((first, last)) = state.pending_hunks.first().zip(state.pending_hunks.last())
@@ -1402,11 +1398,19 @@ impl BufferDiff {
             state.pending_hunks = SumTree::new(buffer);
         }
 
-        cx.emit(BufferDiffEvent::DiffChanged {
-            changed_range: changed_range.clone(),
-            base_text_changed_range,
-        });
-        changed_range
+        cx.spawn(async move |this, cx| {
+            if let Some(parsing_idle) = parsing_idle {
+                parsing_idle.await;
+            }
+            this.update(cx, |_, cx| {
+                cx.emit(BufferDiffEvent::DiffChanged {
+                    changed_range: changed_range.clone(),
+                    base_text_changed_range,
+                });
+            })
+            .ok();
+            changed_range
+        })
     }
 
     pub fn base_text(&self, cx: &App) -> language::BufferSnapshot {
@@ -1454,10 +1458,12 @@ impl BufferDiff {
                 return;
             };
             let state = state.await;
-            this.update(cx, |this, cx| {
-                this.set_snapshot(state, &buffer, cx);
-            })
-            .log_err();
+            if let Some(task) = this
+                .update(cx, |this, cx| this.set_snapshot(state, &buffer, cx))
+                .log_err()
+            {
+                task.await;
+            }
             drop(complete_on_drop)
         })
         .detach();
@@ -1475,8 +1481,9 @@ impl BufferDiff {
         let language = self.base_text(cx).language().cloned();
         let base_text = self.base_text_string(cx).map(|s| s.as_str().into());
         let fut = self.update_diff(buffer.clone(), base_text, false, language, cx);
-        let snapshot = cx.background_executor().block(fut);
-        self.set_snapshot(snapshot, &buffer, cx);
+        let executor = cx.background_executor().clone();
+        let snapshot = executor.block(fut);
+        self.set_snapshot(snapshot, &buffer, cx).detach();
     }
 
     pub fn base_text_buffer(&self) -> Entity<language::Buffer> {
@@ -2589,6 +2596,7 @@ mod tests {
         let diff = cx.new(|cx| {
             BufferDiff::new_with_base_text(&base_text, &buffer.read(cx).text_snapshot(), cx)
         });
+        cx.run_until_parked();
         let (tx, rx) = mpsc::channel();
         let subscription =
             cx.update(|cx| cx.subscribe(&diff, move |_, event, _| tx.send(event.clone()).unwrap()));
@@ -2619,7 +2627,8 @@ mod tests {
                 )
             })
             .await;
-        diff.update(cx, |diff, cx| diff.set_snapshot(update, &snapshot, cx));
+        diff.update(cx, |diff, cx| diff.set_snapshot(update, &snapshot, cx))
+            .await;
         cx.run_until_parked();
         drop(subscription);
         let events = rx.into_iter().collect::<Vec<_>>();

crates/edit_prediction_ui/src/rate_prediction_modal.rs

@@ -335,9 +335,14 @@ impl RatePredictionsModal {
                     );
                     cx.spawn(async move |diff, cx| {
                         let update = update.await;
-                        diff.update(cx, |diff, cx| {
-                            diff.set_snapshot(update, &new_buffer_snapshot.text, cx);
-                        })
+                        if let Some(task) = diff
+                            .update(cx, |diff, cx| {
+                                diff.set_snapshot(update, &new_buffer_snapshot.text, cx)
+                            })
+                            .ok()
+                        {
+                            task.await;
+                        }
                     })
                     .detach();
                 });

crates/git_ui/src/commit_view.rs

@@ -806,8 +806,8 @@ async fn build_buffer_diff(
     diff.update(cx, |diff, cx| {
         diff.language_changed(language, Some(language_registry.clone()), cx);
         diff.set_snapshot(update, &buffer.text, cx)
-    })
-    .ok();
+    })?
+    .await;
 
     Ok(diff)
 }

crates/git_ui/src/file_diff_view.rs

@@ -191,8 +191,9 @@ async fn build_buffer_diff(
             Some(language_registry),
             cx,
         );
-        diff.set_snapshot(update, &new_buffer_snapshot.text, cx);
-    })?;
+        diff.set_snapshot(update, &new_buffer_snapshot.text, cx)
+    })?
+    .await;
 
     Ok(diff)
 }

crates/git_ui/src/text_diff_view.rs

@@ -275,8 +275,9 @@ async fn update_diff_buffer(
         .await;
 
     diff.update(cx, |diff, cx| {
-        diff.set_snapshot(update, &source_buffer_snapshot.text, cx);
-    })?;
+        diff.set_snapshot(update, &source_buffer_snapshot.text, cx)
+    })?
+    .await;
     Ok(())
 }
 

crates/multi_buffer/src/multi_buffer_tests.rs

@@ -4019,8 +4019,9 @@ async fn test_singleton_with_inverted_diff(cx: &mut TestAppContext) {
         })
         .await;
     diff.update(cx, |diff, cx| {
-        diff.set_snapshot(update, &buffer.read(cx).text_snapshot(), cx);
-    });
+        diff.set_snapshot(update, &buffer.read(cx).text_snapshot(), cx)
+    })
+    .await;
     cx.run_until_parked();
 
     assert_new_snapshot(
@@ -4056,8 +4057,9 @@ async fn test_singleton_with_inverted_diff(cx: &mut TestAppContext) {
         })
         .await;
     diff.update(cx, |diff, cx| {
-        diff.set_snapshot(update, &buffer.read(cx).text_snapshot(), cx);
-    });
+        diff.set_snapshot(update, &buffer.read(cx).text_snapshot(), cx)
+    })
+    .await;
     cx.run_until_parked();
 
     assert_new_snapshot(

crates/project/src/git_store.rs

@@ -3159,12 +3159,13 @@ impl BufferGitState {
             let unstaged_changed_range = if let Some((unstaged_diff, new_unstaged_diff)) =
                 unstaged_diff.as_ref().zip(new_unstaged_diff.clone())
             {
-                unstaged_diff.update(cx, |diff, cx| {
+                let task = unstaged_diff.update(cx, |diff, cx| {
                     if language_changed {
                         diff.language_changed(language.clone(), language_registry.clone(), cx);
                     }
                     diff.set_snapshot(new_unstaged_diff, &buffer, cx)
-                })?
+                })?;
+                Some(task.await)
             } else {
                 None
             };
@@ -3174,18 +3175,20 @@ impl BufferGitState {
             if let Some((uncommitted_diff, new_uncommitted_diff)) =
                 uncommitted_diff.as_ref().zip(new_uncommitted_diff.clone())
             {
-                uncommitted_diff.update(cx, |diff, cx| {
-                    if language_changed {
-                        diff.language_changed(language, language_registry, cx);
-                    }
-                    diff.set_snapshot_with_secondary(
-                        new_uncommitted_diff,
-                        &buffer,
-                        unstaged_changed_range,
-                        true,
-                        cx,
-                    );
-                })?;
+                uncommitted_diff
+                    .update(cx, |diff, cx| {
+                        if language_changed {
+                            diff.language_changed(language, language_registry, cx);
+                        }
+                        diff.set_snapshot_with_secondary(
+                            new_uncommitted_diff,
+                            &buffer,
+                            unstaged_changed_range.flatten(),
+                            true,
+                            cx,
+                        )
+                    })?
+                    .await;
             }
 
             log::debug!(