From dd643c949cb415529da310970354c6cc6d53b2eb Mon Sep 17 00:00:00 2001
From: Evan Mesterhazy <emesterhazy@google.com>
Date: Sat, 20 Apr 2024 10:48:52 -0400
Subject: [PATCH] Rewrite rebasing in `jj split` using
 `transform_descendants()`

This is following on the rewrite for `parallelize`.

- https://github.com/martinvonz/jj/pull/3521

Since rebase_descendants from rebase.rs is no longer used outside of that file,
it can be made private again.
---
 cli/src/commands/rebase.rs      |  2 +-
 cli/src/commands/split.rs       | 68 +++++++++------------------------
 cli/tests/test_split_command.rs |  4 +-
 3 files changed, 22 insertions(+), 52 deletions(-)

diff --git a/cli/src/commands/rebase.rs b/cli/src/commands/rebase.rs
index 8b6ffc68cd..bd50dcd243 100644
--- a/cli/src/commands/rebase.rs
+++ b/cli/src/commands/rebase.rs
@@ -286,7 +286,7 @@ fn rebase_branch(
 }
 
 /// Rebases `old_commits` onto `new_parents`.
-pub fn rebase_descendants(
+fn rebase_descendants(
     tx: &mut WorkspaceCommandTransaction,
     settings: &UserSettings,
     new_parents: Vec<Commit>,
diff --git a/cli/src/commands/split.rs b/cli/src/commands/split.rs
index dc5608c669..c4298d23a1 100644
--- a/cli/src/commands/split.rs
+++ b/cli/src/commands/split.rs
@@ -13,18 +13,13 @@
 // limitations under the License.
 use std::io::Write;
 
-use itertools::Itertools;
-use jj_lib::commit::Commit;
 use jj_lib::object_id::ObjectId;
 use jj_lib::repo::Repo;
-use jj_lib::revset::{RevsetExpression, RevsetIteratorExt};
 use jj_lib::rewrite::merge_commit_trees;
-use jj_lib::settings::UserSettings;
 use tracing::instrument;
 
-use crate::cli_util::{CommandHelper, RevisionArg, WorkspaceCommandTransaction};
+use crate::cli_util::{CommandHelper, RevisionArg};
 use crate::command_error::CommandError;
-use crate::commands::rebase::rebase_descendants;
 use crate::description_util::{description_template_for_commit, edit_description};
 use crate::ui::Ui;
 
@@ -179,16 +174,24 @@ the operation will be aborted.
     // commit.
     tx.mut_repo()
         .set_rewritten_commit(commit.id().clone(), second_commit.id().clone());
-    let num_rebased = if args.siblings {
-        rebase_children_for_siblings_split(
-            &mut tx,
-            command.settings(),
-            &commit,
-            vec![first_commit.clone(), second_commit.clone()],
-        )?
-    } else {
-        tx.mut_repo().rebase_descendants(command.settings())?
-    };
+    let mut num_rebased = 0;
+    tx.mut_repo().transform_descendants(
+        command.settings(),
+        vec![commit.id().clone()],
+        |mut rewriter| {
+            num_rebased += 1;
+            if args.siblings {
+                rewriter.replace_parent(
+                    second_commit.id(),
+                    &[first_commit.id().clone(), second_commit.id().clone()],
+                );
+            }
+            // We don't need to do anything special for the non-siblings case
+            // since we already marked the original commit as rewritten.
+            rewriter.rebase(command.settings())?.write()?;
+            Ok(())
+        },
+    )?;
 
     if let Some(mut formatter) = ui.status_formatter() {
         if num_rebased > 0 {
@@ -203,36 +206,3 @@ the operation will be aborted.
     tx.finish(ui, format!("split commit {}", commit.id().hex()))?;
     Ok(())
 }
-
-// Rebases the children of `original_commit` by replacing `original_commit` with
-// `new_siblings`. Any parents other than `original_commit` will remain after
-// the rebase.
-fn rebase_children_for_siblings_split(
-    tx: &mut WorkspaceCommandTransaction,
-    settings: &UserSettings,
-    original_commit: &Commit,
-    new_siblings: Vec<Commit>,
-) -> Result<usize, CommandError> {
-    let children: Vec<Commit> = RevsetExpression::commit(original_commit.id().clone())
-        .children()
-        .evaluate_programmatic(tx.base_repo().as_ref())?
-        .iter()
-        .commits(tx.base_repo().store())
-        .try_collect()?;
-    let mut num_rebased = 0;
-    for child in children {
-        let new_parents = child
-            .parents()
-            .into_iter()
-            .flat_map(|c| {
-                if c.id() == original_commit.id() {
-                    new_siblings.clone()
-                } else {
-                    vec![c]
-                }
-            })
-            .collect_vec();
-        num_rebased += rebase_descendants(tx, settings, new_parents, &[child], Default::default())?;
-    }
-    Ok(num_rebased)
-}
diff --git a/cli/tests/test_split_command.rs b/cli/tests/test_split_command.rs
index 0d961e1df1..051fe24f72 100644
--- a/cli/tests/test_split_command.rs
+++ b/cli/tests/test_split_command.rs
@@ -346,8 +346,8 @@ fn test_split_siblings_no_descendants() {
     let (stdout, stderr) = test_env.jj_cmd_ok(&workspace_path, &["split", "--siblings", "file1"]);
     insta::assert_snapshot!(stdout, @"");
     insta::assert_snapshot!(stderr, @r###"
-    First part: qpvuntsm?? 8d2b7558 TESTED=TODO
-    Second part: zsuskuln acd41528 (no description set)
+    First part: qpvuntsm 8d2b7558 TESTED=TODO
+    Second part: zsuskuln acd41528 test_branch | (no description set)
     Working copy now at: zsuskuln acd41528 test_branch | (no description set)
     Parent commit      : zzzzzzzz 00000000 (empty) (no description set)
     Added 0 files, modified 0 files, removed 1 files