repo: add method for tranforming descendants, use in rebase_descendants()

There are several existing commands that would benefit from an API
that makes it easier to rewrite a whole graph of commits while
transforming them in some way.

`jj squash` is one example. When squashing into an ancestor, that
command currently rewrites the ancestor, then rebases descendants, and
then rewrites the rewritten source commit. It would be better to
rewrite the source commit (and any descendants) only once.

Another example is the future `jj fix`. That command will want to
rewrite a graph while updating the trees. There's currently no good
API for that; you have to manually iterate over descendants and
rewrite them.

This patch adds a new `MutableRepo::transform_descendants()` method
that takes a callback which gets a `CommitRewriter` passed to it. The
callback can then decide to change the parents, the tree, etc. The
callback is also free to leave the commit in place or to abandon it.

I updated the regular `rebase_descendants()` to use the new function
in order to exercise it. I hope we can replace all of the
`rebase_descendant_*()` flavors later.

I added a `replace_parent()` method that was a bit useful for the test
case. It could easily be hard-coded in the test case instead, but I
think the method will be useful for `jj git sync` and similar in the
future.

lib/src/repo.rs

@@ -51,7 +51,7 @@ use crate::refs::{
     diff_named_ref_targets, diff_named_remote_refs, merge_ref_targets, merge_remote_refs,
 };
 use crate::revset::{RevsetEvaluationError, RevsetExpression, RevsetIteratorExt};
-use crate::rewrite::{DescendantRebaser, RebaseOptions};
+use crate::rewrite::{CommitRewriter, DescendantRebaser, RebaseOptions};
 use crate::settings::{RepoSettings, UserSettings};
 use crate::signing::{SignInitError, Signer};
 use crate::simple_op_heads_store::SimpleOpHeadsStore;
@@ -1154,6 +1154,42 @@ impl MutableRepo {
         )
     }
 
+    /// Rewrite descendants of the given roots.
+    ///
+    /// The callback will be called for each commit with the new parents
+    /// prepopulated. The callback may change the parents and write the new
+    /// commit, or it may abandon the commit, or it may leave the old commit
+    /// unchanged.
+    ///
+    /// The set of commits to visit is determined at the start. If the callback
+    /// adds new descendants, then the callback will not be called for those.
+    /// Similarly, if the callback rewrites unrelated commits, then the callback
+    /// will not be called for descendants of those commits.
+    pub fn transform_descendants(
+        &mut self,
+        settings: &UserSettings,
+        roots: Vec<CommitId>,
+        mut callback: impl FnMut(CommitRewriter) -> BackendResult<()>,
+    ) -> BackendResult<()> {
+        let mut to_visit = self.find_descendants_to_rebase(roots)?;
+        while let Some(old_commit) = to_visit.pop() {
+            let new_parent_ids = self.new_parents(old_commit.parent_ids());
+            let rewriter = CommitRewriter::new(self, old_commit, new_parent_ids);
+            callback(rewriter)?;
+        }
+        self.update_rewritten_references(settings)?;
+        // Since we didn't necessarily visit all descendants of rewritten commits (e.g.
+        // if they were rewritten in the callback), there can still be commits left to
+        // rebase, so we don't clear `parent_mapping` here.
+        // TODO: Should we make this stricter? We could check that there were no
+        // rewrites before this function was called, and we can check that only
+        // commits in the `to_visit` set were added by the callback. Then we
+        // could clear `parent_mapping` here and not have to scan it again at
+        // the end of the transaction when we call `rebase_descendants()`.
+
+        Ok(())
+    }
+
     /// After the rebaser returned by this function is dropped,
     /// self.parent_mapping needs to be cleared.
     fn rebase_descendants_return_rebaser<'settings, 'repo>(
@@ -1202,6 +1238,7 @@ impl MutableRepo {
     /// **its parent**. One can tell this case apart since the change ids of the
     /// key and the value will not match. The parent will inherit the
     /// descendants and the branches of the abandoned commit.
+    // TODO: Rewrite this using `transform_descendants()`
     pub fn rebase_descendants_with_options_return_map(
         &mut self,
         settings: &UserSettings,
@@ -1218,7 +1255,18 @@ impl MutableRepo {
     }
 
     pub fn rebase_descendants(&mut self, settings: &UserSettings) -> BackendResult<usize> {
-        self.rebase_descendants_with_options(settings, Default::default())
+        let roots = self.parent_mapping.keys().cloned().collect_vec();
+        let mut num_rebased = 0;
+        self.transform_descendants(settings, roots, |rewriter| {
+            if rewriter.parents_changed() {
+                let builder = rewriter.rebase(settings)?;
+                builder.write()?;
+                num_rebased += 1;
+            }
+            Ok(())
+        })?;
+        self.parent_mapping.clear();
+        Ok(num_rebased)
     }
 
     pub fn rebase_descendants_return_map(

lib/src/rewrite.rs

@@ -145,6 +145,15 @@ impl<'repo> CommitRewriter<'repo> {
         self.new_parents = new_parents;
     }
 
+    /// Update the intended new parents by replacing any occurrence of
+    /// `old_parent` by `new_parents`
+    pub fn replace_parent(&mut self, old_parent: &CommitId, new_parents: &[CommitId]) {
+        if let Some(i) = self.new_parents.iter().position(|p| p == old_parent) {
+            self.new_parents
+                .splice(i..i + 1, new_parents.iter().cloned());
+        }
+    }
+
     /// Checks if the intended new parents are different from the old commit's
     /// parents.
     pub fn parents_changed(&self) -> bool {

lib/tests/runner.rs

@@ -29,6 +29,7 @@ mod test_operations;
 mod test_refs;
 mod test_revset;
 mod test_rewrite;
+mod test_rewrite_transform;
 mod test_signing;
 mod test_ssh_signing;
 mod test_view;

lib/tests/test_rewrite_transform.rs

@@ -0,0 +1,84 @@
+// Copyright 2024 The Jujutsu Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::HashMap;
+
+use jj_lib::repo::Repo;
+use maplit::hashset;
+use testutils::{CommitGraphBuilder, TestRepo};
+
+// Simulate some `jj sync` command that rebases B:: onto G while abandoning C
+// (because it's presumably already in G).
+//
+// G
+// | E
+// | D F
+// | |/
+// | C
+// | B
+// |/
+// A
+#[test]
+fn test_transform_descendants_sync() {
+    let settings = testutils::user_settings();
+    let test_repo = TestRepo::init();
+    let repo = &test_repo.repo;
+
+    let mut tx = repo.start_transaction(&settings);
+    let mut graph_builder = CommitGraphBuilder::new(&settings, tx.mut_repo());
+    let commit_a = graph_builder.initial_commit();
+    let commit_b = graph_builder.commit_with_parents(&[&commit_a]);
+    let commit_c = graph_builder.commit_with_parents(&[&commit_b]);
+    let commit_d = graph_builder.commit_with_parents(&[&commit_c]);
+    let commit_e = graph_builder.commit_with_parents(&[&commit_d]);
+    let commit_f = graph_builder.commit_with_parents(&[&commit_c]);
+    let commit_g = graph_builder.commit_with_parents(&[&commit_a]);
+
+    let mut rebased = HashMap::new();
+    tx.mut_repo()
+        .transform_descendants(&settings, vec![commit_b.id().clone()], |mut rewriter| {
+            rewriter.replace_parent(commit_a.id(), &[commit_g.id().clone()]);
+            if *rewriter.old_commit() == commit_c {
+                let old_id = rewriter.old_commit().id().clone();
+                let new_parent_ids = rewriter.new_parents().to_vec();
+                rewriter
+                    .mut_repo()
+                    .record_abandoned_commit_with_parents(old_id, new_parent_ids);
+            } else {
+                let old_commit_id = rewriter.old_commit().id().clone();
+                let new_commit = rewriter.rebase(&settings)?.write()?;
+                rebased.insert(old_commit_id, new_commit);
+            }
+            Ok(())
+        })
+        .unwrap();
+    assert_eq!(rebased.len(), 4);
+    let new_commit_b = rebased.get(commit_b.id()).unwrap();
+    let new_commit_d = rebased.get(commit_d.id()).unwrap();
+    let new_commit_e = rebased.get(commit_e.id()).unwrap();
+    let new_commit_f = rebased.get(commit_f.id()).unwrap();
+
+    assert_eq!(
+        *tx.mut_repo().view().heads(),
+        hashset! {
+            new_commit_e.id().clone(),
+            new_commit_f.id().clone(),
+        }
+    );
+
+    assert_eq!(new_commit_b.parent_ids(), vec![commit_g.id().clone()]);
+    assert_eq!(new_commit_d.parent_ids(), vec![new_commit_b.id().clone()]);
+    assert_eq!(new_commit_e.parent_ids(), vec![new_commit_d.id().clone()]);
+    assert_eq!(new_commit_f.parent_ids(), vec![new_commit_b.id().clone()]);
+}