feat: follow forward working-checkout changes in diffs

src/repository.rs

@@ -307,6 +307,10 @@ struct PreparedComparison {
     left: Snapshot,
     right: Snapshot,
     files: Vec<(PathBuf, Option<PathBuf>, FileStatus)>,
+    before: String,
+    after: String,
+    working_after: bool,
+    working_anchor: Option<String>,
 }
 
 pub struct Repository {
@@ -1331,6 +1335,53 @@ impl Repository {
         }
         Ok(base.into())
     }
+    fn working_anchor(&self) -> Result<Option<String>, String> {
+        if self.kind == RepoKind::Jujutsu {
+            return Ok(Some(
+                text(self.jj(
+                    &self.root,
+                    &["log", "--no-graph", "-r", "@", "-T", "change_id"],
+                )?)?
+                .trim()
+                .into(),
+            ));
+        }
+        let mut command = git_command(&self.store);
+        command.args(["rev-parse", "--verify", "--quiet", "HEAD^{commit}"]);
+        let output = run(command, vec![], &self.cancelled)?;
+        if output.code == Some(1) {
+            return Ok(None); // An unborn checkout has no commit anchor yet.
+        }
+        Ok(Some(text(checked(output)?)?.trim().into()))
+    }
+    fn working_advanced(&self, before: Option<&str>, after: Option<&str>) -> Result<bool, String> {
+        if before == after || after.is_none() {
+            return Ok(false);
+        }
+        let Some(before) = before else {
+            return Ok(true);
+        };
+        if self.kind == RepoKind::Jujutsu {
+            let revset = format!(
+                "ancestors(@) & change_id({})",
+                serde_json::to_string(before).unwrap()
+            );
+            return Ok(!self
+                .jj(
+                    &self.root,
+                    &["log", "--no-graph", "-r", &revset, "-T", "change_id"],
+                )?
+                .is_empty());
+        }
+        let mut command = git_command(&self.store);
+        command.args(["merge-base", "--is-ancestor", before, after.unwrap()]);
+        let output = run(command, vec![], &self.cancelled)?;
+        if output.code == Some(1) {
+            Ok(false)
+        } else {
+            checked(output).map(|_| true)
+        }
+    }
     /// Freeze comparison sources and return path/status metadata only.
     /// A new comparison replaces sources retained for the previous request ID.
     pub fn compare(
@@ -1341,7 +1392,7 @@ impl Repository {
         mode: ComparisonMode,
         buffers: &[(PathBuf, String)],
     ) -> Result<Comparison, String> {
-        self.comparison = None;
+        self.check_cancelled()?;
         let resolved = if mode == ComparisonMode::BranchChanges
             && before.target != RevisionTarget::Working
             && after.target != RevisionTarget::Working
@@ -1365,26 +1416,98 @@ impl Repository {
             (&before.target, &after.target, before.label.as_str()),
             |(before, after, label)| (before, after, label.as_str()),
         );
-        let mut left = self.snapshot(before_target)?;
         let working_after = *after_target == RevisionTarget::Working;
+        let working_anchor = if working_after {
+            self.working_anchor()?
+        } else {
+            None
+        };
+        let left = self.snapshot(before_target)?;
         let same = before_target == after_target && (!working_after || buffers.is_empty());
-        let mut right = if same {
+        let right = if same {
             BTreeMap::new()
         } else if working_after {
             self.working_snapshot(buffers)?
         } else {
             self.snapshot(after_target)?
         };
+        let mut prepared = PreparedComparison {
+            id: comparison_id,
+            left,
+            right,
+            files: vec![],
+            before: before_label.into(),
+            after: after.label.clone(),
+            working_after,
+            working_anchor,
+        };
+        let comparison = self.build_comparison(&mut prepared, same)?;
+        if working_after && self.working_anchor()? != prepared.working_anchor {
+            return Err("Working checkout moved during comparison".into());
+        }
+        self.comparison = Some(prepared);
+        Ok(comparison)
+    }
+
+    /// Advance only a selected Working After, keeping its original Before frozen.
+    pub fn refresh_working_after(
+        &mut self,
+        comparison_id: u64,
+        new_id: u64,
+        buffers: &[(PathBuf, String)],
+    ) -> Result<Option<Comparison>, String> {
+        let Some(prepared) = self
+            .comparison
+            .as_ref()
+            .filter(|prepared| prepared.id == comparison_id && prepared.working_after)
+        else {
+            return Ok(None);
+        };
+        let anchor = self.working_anchor()?;
+        if !self.working_advanced(prepared.working_anchor.as_deref(), anchor.as_deref())? {
+            return Ok(None);
+        }
+        let right = self.working_snapshot(buffers)?;
+        let mut prepared = self.comparison.take().unwrap();
+        let previous_right = std::mem::replace(&mut prepared.right, right);
+        let previous_files = std::mem::take(&mut prepared.files);
+        let result = self
+            .build_comparison(&mut prepared, false)
+            .and_then(|comparison| {
+                self.check_cancelled()?;
+                if self.working_anchor()? != anchor {
+                    return Err("Working checkout moved during comparison".into());
+                }
+                Ok(comparison)
+            });
+        if result.is_ok() {
+            prepared.id = new_id;
+            prepared.working_anchor = anchor;
+        } else {
+            prepared.right = previous_right;
+            prepared.files = previous_files;
+        }
+        self.comparison = Some(prepared);
+        result.map(Some)
+    }
+
+    fn build_comparison(
+        &self,
+        prepared: &mut PreparedComparison,
+        same: bool,
+    ) -> Result<Comparison, String> {
+        let left = &mut prepared.left;
+        let right = &mut prepared.right;
         // Frozen checkout/conflict bytes have no object identity. Compare them in a
         // single batch, while immutable tree/tree comparisons need only object IDs.
         let frozen = |file: &FileData| {
             file.object.is_none() && matches!(file.mode.as_str(), "100644" | "100755" | "120000")
         };
         if !same && (left.values().any(frozen) || right.values().any(frozen)) {
-            self.materialize(&mut left)?;
-            self.materialize(&mut right)?;
+            self.materialize(left)?;
+            self.materialize(right)?;
         }
-        let right_view = if same { &left } else { &right };
+        let right_view = if same { &*left } else { &*right };
         let mut files = vec![];
         let paths: BTreeSet<_> = left.keys().chain(right_view.keys()).collect();
         for path in paths {
@@ -1442,7 +1565,7 @@ impl Repository {
             .enumerate()
             .filter_map(|(i, f)| (!removed.contains(&i)).then_some(f))
             .collect();
-        let editable_paths = if working_after {
+        let editable_paths = if prepared.working_after {
             files
                 .iter()
                 .filter(|file| {
@@ -1455,18 +1578,13 @@ impl Repository {
         } else {
             BTreeSet::new()
         };
-        self.comparison = Some(PreparedComparison {
-            id: comparison_id,
-            left,
-            right,
-            files: files
-                .iter()
-                .map(|file| (file.path.clone(), file.old_path.clone(), file.status))
-                .collect(),
-        });
+        prepared.files = files
+            .iter()
+            .map(|file| (file.path.clone(), file.old_path.clone(), file.status))
+            .collect();
         Ok(Comparison {
-            before: before_label.into(),
-            after: after.label.clone(),
+            before: prepared.before.clone(),
+            after: prepared.after.clone(),
             files,
             root: self.root.clone(),
             editable_paths,