Fix missing hunks in project diff after revert (#25906)

Release Notes:

- N/A

Cargo.lock

@@ -5420,8 +5420,10 @@ dependencies = [
  "buffer_diff",
  "collections",
  "component",
+ "ctor",
  "db",
  "editor",
+ "env_logger 0.11.6",
  "feature_flags",
  "futures 0.3.31",
  "fuzzy",

crates/git_ui/Cargo.toml

@@ -57,6 +57,8 @@ zed_actions.workspace = true
 windows.workspace = true
 
 [dev-dependencies]
+ctor.workspace = true
+env_logger.workspace = true
 editor = { workspace = true, features = ["test-support"] }
 gpui = { workspace = true, features = ["test-support"] }
 project = { workspace = true, features = ["test-support"] }

crates/git_ui/src/project_diff.rs

@@ -400,6 +400,7 @@ impl ProjectDiff {
         self.editor.update(cx, |editor, cx| {
             if was_empty {
                 editor.change_selections(None, window, cx, |selections| {
+                    // TODO select the very beginning (possibly inside a deletion)
                     selections.select_ranges([0..0])
                 });
             }
@@ -980,6 +981,11 @@ mod tests {
 
     use super::*;
 
+    #[ctor::ctor]
+    fn init_logger() {
+        env_logger::init();
+    }
+
     fn init_test(cx: &mut TestAppContext) {
         cx.update(|cx| {
             let store = SettingsStore::test(cx);
@@ -1152,4 +1158,107 @@ mod tests {
             .unindent(),
         );
     }
+
+    #[gpui::test]
+    async fn test_hunks_after_restore_then_modify(cx: &mut TestAppContext) {
+        init_test(cx);
+
+        let fs = FakeFs::new(cx.executor());
+        fs.insert_tree(
+            path!("/project"),
+            json!({
+                ".git": {},
+                "foo": "modified\n",
+            }),
+        )
+        .await;
+        let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
+        let (workspace, cx) =
+            cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx));
+        let buffer = project
+            .update(cx, |project, cx| {
+                project.open_local_buffer(path!("/project/foo"), cx)
+            })
+            .await
+            .unwrap();
+        let buffer_editor = cx.new_window_entity(|window, cx| {
+            Editor::for_buffer(buffer, Some(project.clone()), window, cx)
+        });
+        let diff = cx.new_window_entity(|window, cx| {
+            ProjectDiff::new(project.clone(), workspace, window, cx)
+        });
+        cx.run_until_parked();
+
+        fs.set_head_for_repo(
+            path!("/project/.git").as_ref(),
+            &[("foo".into(), "original\n".into())],
+        );
+        fs.with_git_state(path!("/project/.git").as_ref(), true, |state| {
+            state.statuses = HashMap::from_iter([(
+                "foo".into(),
+                TrackedStatus {
+                    index_status: StatusCode::Unmodified,
+                    worktree_status: StatusCode::Modified,
+                }
+                .into(),
+            )]);
+        });
+        cx.run_until_parked();
+
+        let diff_editor = diff.update(cx, |diff, _| diff.editor.clone());
+
+        assert_state_with_diff(
+            &diff_editor,
+            cx,
+            &"
+                - original
+                + ˇmodified
+            "
+            .unindent(),
+        );
+
+        let prev_buffer_hunks =
+            cx.update_window_entity(&buffer_editor, |buffer_editor, window, cx| {
+                let snapshot = buffer_editor.snapshot(window, cx);
+                let snapshot = &snapshot.buffer_snapshot;
+                let prev_buffer_hunks = buffer_editor
+                    .diff_hunks_in_ranges(&[editor::Anchor::min()..editor::Anchor::max()], snapshot)
+                    .collect::<Vec<_>>();
+                buffer_editor.git_restore(&Default::default(), window, cx);
+                prev_buffer_hunks
+            });
+        assert_eq!(prev_buffer_hunks.len(), 1);
+        cx.run_until_parked();
+
+        let new_buffer_hunks =
+            cx.update_window_entity(&buffer_editor, |buffer_editor, window, cx| {
+                let snapshot = buffer_editor.snapshot(window, cx);
+                let snapshot = &snapshot.buffer_snapshot;
+                let new_buffer_hunks = buffer_editor
+                    .diff_hunks_in_ranges(&[editor::Anchor::min()..editor::Anchor::max()], snapshot)
+                    .collect::<Vec<_>>();
+                buffer_editor.git_restore(&Default::default(), window, cx);
+                new_buffer_hunks
+            });
+        assert_eq!(new_buffer_hunks.as_slice(), &[]);
+
+        cx.update_window_entity(&buffer_editor, |buffer_editor, window, cx| {
+            buffer_editor.set_text("different\n", window, cx);
+            buffer_editor.save(false, project.clone(), window, cx)
+        })
+        .await
+        .unwrap();
+
+        cx.run_until_parked();
+
+        assert_state_with_diff(
+            &diff_editor,
+            cx,
+            &"
+                - original
+                + ˇdifferent
+            "
+            .unindent(),
+        );
+    }
 }

crates/multi_buffer/src/multi_buffer.rs

@@ -2062,7 +2062,12 @@ impl MultiBuffer {
                     if let Some(buffer_state) = buffers.get_mut(&excerpt.buffer_id) {
                         buffer_state.excerpts.retain(|l| l != &excerpt.locator);
                         if buffer_state.excerpts.is_empty() {
+                            log::debug!(
+                                "removing buffer and diff for buffer {}",
+                                excerpt.buffer_id
+                            );
                             buffers.remove(&excerpt.buffer_id);
+                            self.diffs.remove(&excerpt.buffer_id);
                         }
                     }
                     cursor.next(&());
@@ -2716,6 +2721,11 @@ impl MultiBuffer {
         snapshot.has_deleted_file = has_deleted_file;
         snapshot.has_conflict = has_conflict;
 
+        snapshot.diffs.retain(|_, _| false);
+        for (id, diff) in self.diffs.iter() {
+            snapshot.diffs.insert(*id, diff.diff.read(cx).snapshot(cx));
+        }
+
         excerpts_to_edit.sort_unstable_by_key(|(locator, _, _)| *locator);
 
         let mut edits = Vec::new();
@@ -3476,7 +3486,10 @@ impl MultiBufferSnapshot {
     ) -> impl Iterator<Item = MultiBufferDiffHunk> + '_ {
         let query_range = range.start.to_point(self)..range.end.to_point(self);
         self.lift_buffer_metadata(query_range.clone(), move |buffer, buffer_range| {
-            let diff = self.diffs.get(&buffer.remote_id())?;
+            let Some(diff) = self.diffs.get(&buffer.remote_id()) else {
+                log::debug!("no diff found for {:?}", buffer.remote_id());
+                return None;
+            };
             let buffer_start = buffer.anchor_before(buffer_range.start);
             let buffer_end = buffer.anchor_after(buffer_range.end);
             Some(

crates/multi_buffer/src/multi_buffer_tests.rs

@@ -2160,6 +2160,7 @@ impl ReferenceMultibuffer {
             .unwrap();
         let excerpt = self.excerpts.remove(ix);
         let buffer = excerpt.buffer.read(cx);
+        let id = buffer.remote_id();
         log::info!(
             "Removing excerpt {}: {:?}",
             ix,
@@ -2167,6 +2168,13 @@ impl ReferenceMultibuffer {
                 .text_for_range(excerpt.range.to_offset(buffer))
                 .collect::<String>(),
         );
+        if !self
+            .excerpts
+            .iter()
+            .any(|excerpt| excerpt.buffer.read(cx).remote_id() == id)
+        {
+            self.diffs.remove(&id);
+        }
     }
 
     fn insert_excerpt_after(
@@ -2415,6 +2423,7 @@ async fn test_random_multibuffer(cx: &mut TestAppContext, mut rng: StdRng) {
         .unwrap_or(10);
 
     let mut buffers: Vec<Entity<Buffer>> = Vec::new();
+    let mut base_texts: HashMap<BufferId, String> = HashMap::default();
     let multibuffer = cx.new(|_| MultiBuffer::new(Capability::ReadWrite));
     let mut reference = ReferenceMultibuffer::default();
     let mut anchors = Vec::new();
@@ -2522,9 +2531,10 @@ async fn test_random_multibuffer(cx: &mut TestAppContext, mut rng: StdRng) {
                         ..snapshot.anchor_in_excerpt(excerpt.id, end).unwrap();
 
                     log::info!(
-                        "expanding diff hunks in range {:?} (excerpt id {:?}) index {excerpt_ix:?})",
+                        "expanding diff hunks in range {:?} (excerpt id {:?}, index {excerpt_ix:?}, buffer id {:?})",
                         range.to_offset(&snapshot),
-                        excerpt.id
+                        excerpt.id,
+                        excerpt.buffer.read(cx).remote_id(),
                     );
                     reference.expand_diff_hunks(excerpt.id, start..end, cx);
                     multibuffer.expand_diff_hunks(vec![range], cx);
@@ -2534,7 +2544,7 @@ async fn test_random_multibuffer(cx: &mut TestAppContext, mut rng: StdRng) {
                 multibuffer.update(cx, |multibuffer, cx| {
                     for buffer in multibuffer.all_buffers() {
                         let snapshot = buffer.read(cx).snapshot();
-                        let _ = multibuffer.diff_for(snapshot.remote_id()).unwrap().update(
+                        multibuffer.diff_for(snapshot.remote_id()).unwrap().update(
                             cx,
                             |diff, cx| {
                                 log::info!(
@@ -2551,17 +2561,16 @@ async fn test_random_multibuffer(cx: &mut TestAppContext, mut rng: StdRng) {
             }
             _ => {
                 let buffer_handle = if buffers.is_empty() || rng.gen_bool(0.4) {
-                    let base_text = util::RandomCharIter::new(&mut rng)
+                    let mut base_text = util::RandomCharIter::new(&mut rng)
                         .take(256)
                         .collect::<String>();
 
                     let buffer = cx.new(|cx| Buffer::local(base_text.clone(), cx));
-                    let diff = cx.new(|cx| BufferDiff::new_with_base_text(&base_text, &buffer, cx));
+                    text::LineEnding::normalize(&mut base_text);
-
+                    base_texts.insert(
-                    multibuffer.update(cx, |multibuffer, cx| {
+                        buffer.read_with(cx, |buffer, _| buffer.remote_id()),
-                        reference.add_diff(diff.clone(), cx);
+                        base_text,
-                        multibuffer.add_diff(diff, cx)
+                    );
-                    });
                     buffers.push(buffer);
                     buffers.last().unwrap()
                 } else {
@@ -2595,6 +2604,18 @@ async fn test_random_multibuffer(cx: &mut TestAppContext, mut rng: StdRng) {
                     (start_ix..end_ix, anchor_range)
                 });
 
+                multibuffer.update(cx, |multibuffer, cx| {
+                    let id = buffer_handle.read(cx).remote_id();
+                    if multibuffer.diff_for(id).is_none() {
+                        let base_text = base_texts.get(&id).unwrap();
+                        let diff = cx.new(|cx| {
+                            BufferDiff::new_with_base_text(base_text, &buffer_handle, cx)
+                        });
+                        reference.add_diff(diff.clone(), cx);
+                        multibuffer.add_diff(diff, cx)
+                    }
+                });
+
                 let excerpt_id = multibuffer.update(cx, |multibuffer, cx| {
                     multibuffer
                         .insert_excerpts_after(