cli rebase: Allow jj rebase -r to rebase a commit onto a descendant

CHANGELOG.md

@@ -72,6 +72,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * `jj workspace add` now preserves all parents of the old working-copy commit
   instead of just the first one.
 
+* `jj rebase -r` gained the ability to rebase a revision `A` onto a descendant
+  of `A`.
+
 ### Fixed bugs
 
 * Updating the working copy to a commit where a file that's currently ignored

cli/src/commands/rebase.rs

@@ -11,6 +11,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
 use std::io::Write;
 use std::sync::Arc;
 
@@ -137,6 +138,9 @@ pub(crate) struct RebaseArgs {
     revision: Option<RevisionArg>,
     /// The revision(s) to rebase onto (can be repeated to create a merge
     /// commit)
+    ///
+    /// Unlike `-s` or `-b`, you may `jj rebase -r` a revision `A` onto a
+    /// descendant of `A`.
     #[arg(long, short, required = true)]
     destination: Vec<RevisionArg>,
     /// Deprecated. Please prefix the revset with `all:` instead.
@@ -264,7 +268,13 @@ fn rebase_revision(
 ) -> Result<(), CommandError> {
     let old_commit = workspace_command.resolve_single_rev(rev_str, ui)?;
     workspace_command.check_rewritable([&old_commit])?;
-    check_rebase_destinations(workspace_command.repo(), new_parents, &old_commit)?;
+    if new_parents.contains(&old_commit) {
+        return Err(user_error(format!(
+            "Cannot rebase {} onto itself",
+            short_commit_hash(old_commit.id()),
+        )));
+    }
+
     let children_expression = RevsetExpression::commit(old_commit.id().clone()).children();
     let child_commits: Vec<_> = children_expression
         .resolve(workspace_command.repo().as_ref())
@@ -274,14 +284,14 @@ fn rebase_revision(
         .iter()
         .commits(workspace_command.repo().store())
         .try_collect()?;
+    // Currently, immutable commits are defied so that a child of a rewriteable
+    // commit is always rewriteable.
+    debug_assert!(workspace_command.check_rewritable(&child_commits).is_ok());
 
+    // First, rebase the children of `old_commit`.
     let mut tx =
         workspace_command.start_transaction(&format!("rebase commit {}", old_commit.id().hex()));
-    rebase_commit(settings, tx.mut_repo(), &old_commit, new_parents)?;
-    // Manually rebase children because we don't want to rebase them onto the
-    // rewritten commit. (But we still want to record the commit as rewritten so
-    // branches and the working copy get updated to the rewritten commit.)
-    let mut num_rebased_descendants = 0;
+    let mut rebased_commit_ids = HashMap::new();
     for child_commit in &child_commits {
         let new_child_parent_ids: Vec<CommitId> = child_commit
             .parents()
@@ -316,10 +326,43 @@ fn rebase_revision(
             .commits(tx.base_repo().store())
             .try_collect()?;
 
-        rebase_commit(settings, tx.mut_repo(), child_commit, &new_child_parents)?;
-        num_rebased_descendants += 1;
+        rebased_commit_ids.insert(
+            child_commit.id().clone(),
+            rebase_commit(settings, tx.mut_repo(), child_commit, &new_child_parents)?
+                .id()
+                .clone(),
+        );
     }
-    num_rebased_descendants += tx.mut_repo().rebase_descendants(settings)?;
+    // Now, rebase the descendants of the children.
+    rebased_commit_ids.extend(tx.mut_repo().rebase_descendants_return_map(settings)?);
+    let num_rebased_descendants = rebased_commit_ids.len();
+
+    // We now update `new_parents` to account for the rebase of all of
+    // `old_commit`'s descendants. Even if some of the original `new_parents` were
+    // descendants of `old_commit`, this will no longer be the case after the
+    // update.
+    //
+    // To make the update simpler, we assume that each commit was rewritten only
+    // once; we don't have a situation where both `(A,B)` and `(B,C)` are in
+    // `rebased_commit_ids`. This assumption relies on the fact that a descendant of
+    // a child of `old_commit` cannot also be a direct child of `old_commit`.
+    let new_parents: Vec<_> = new_parents
+        .iter()
+        .map(|new_parent| {
+            rebased_commit_ids
+                .get(new_parent.id())
+                .map_or(Ok(new_parent.clone()), |rebased_new_parent_id| {
+                    tx.repo().store().get_commit(rebased_new_parent_id)
+                })
+        })
+        .try_collect()?;
+
+    // Finally, it's safe to rebase `old_commit`. At this point, it should no longer
+    // have any children; they have all been rebased and the originals have been
+    // abandoned.
+    rebase_commit(settings, tx.mut_repo(), &old_commit, &new_parents)?;
+    debug_assert_eq!(tx.mut_repo().rebase_descendants(settings)?, 0);
+
     if num_rebased_descendants > 0 {
         writeln!(
             ui.stderr(),

cli/tests/test_rebase_command.rs

@@ -72,10 +72,16 @@ fn test_rebase_invalid() {
     For more information, try '--help'.
     "###);
 
-    // Rebase onto descendant with -r
-    let stderr = test_env.jj_cmd_failure(&repo_path, &["rebase", "-r", "a", "-d", "b"]);
+    // Rebase onto self with -r
+    let stderr = test_env.jj_cmd_failure(&repo_path, &["rebase", "-r", "a", "-d", "a"]);
     insta::assert_snapshot!(stderr, @r###"
-    Error: Cannot rebase 2443ea76b0b1 onto descendant 1394f625cbbd
+    Error: Cannot rebase 2443ea76b0b1 onto itself
+    "###);
+
+    // Rebase root with -r
+    let stderr = test_env.jj_cmd_failure(&repo_path, &["rebase", "-r", "root()", "-d", "a"]);
+    insta::assert_snapshot!(stderr, @r###"
+    Error: The root commit 000000000000 is immutable
     "###);
 
     // Rebase onto descendant with -s
@@ -270,9 +276,9 @@ fn test_rebase_single_revision() {
     Added 0 files, modified 0 files, removed 1 files
     "###);
     insta::assert_snapshot!(get_log_output(&test_env, &repo_path), @r###"
-    @  d
-    ◉  c
-    │ ◉  b
+    ◉  b
+    │ @  d
+    │ ◉  c
     ├─╯
     ◉  a
     ◉
@@ -291,12 +297,12 @@ fn test_rebase_single_revision() {
     Added 0 files, modified 0 files, removed 1 files
     "###);
     insta::assert_snapshot!(get_log_output(&test_env, &repo_path), @r###"
-    @    d
-    ├─╮
-    │ ◉  a
-    ◉ │  b
-    ├─╯
-    │ ◉  c
+    ◉  c
+    │ @    d
+    │ ├─╮
+    │ │ ◉  a
+    ├───╯
+    │ ◉  b
     ├─╯
     ◉
     "###);
@@ -335,15 +341,88 @@ fn test_rebase_single_revision_merge_parent() {
     Added 0 files, modified 0 files, removed 1 files
     "###);
     insta::assert_snapshot!(get_log_output(&test_env, &repo_path), @r###"
-    @    d
+    ◉  c
+    │ @  d
+    ╭─┤
+    ◉ │  a
+    │ ◉  b
+    ├─╯
+    ◉
+    "###);
+}
+
+#[test]
+fn test_rebase_revision_onto_descendant() {
+    let test_env = TestEnvironment::default();
+    test_env.jj_cmd_ok(test_env.env_root(), &["init", "repo", "--git"]);
+    let repo_path = test_env.env_root().join("repo");
+
+    create_commit(&test_env, &repo_path, "base", &[]);
+    create_commit(&test_env, &repo_path, "a", &["base"]);
+    create_commit(&test_env, &repo_path, "b", &["base"]);
+    create_commit(&test_env, &repo_path, "merge", &["b", "a"]);
+    // Test the setup
+    insta::assert_snapshot!(get_log_output(&test_env, &repo_path), @r###"
+    @    merge
     ├─╮
+    │ ◉  a
     ◉ │  b
-    │ │ ◉  c
-    │ ├─╯
+    ├─╯
+    ◉  base
+    ◉
+    "###);
+    let setup_opid = test_env.current_operation_id(&repo_path);
+
+    // Simpler example
+    let (stdout, stderr) = test_env.jj_cmd_ok(&repo_path, &["rebase", "-r", "base", "-d", "a"]);
+    insta::assert_snapshot!(stdout, @"");
+    insta::assert_snapshot!(stderr, @r###"
+    Also rebased 3 descendant commits onto parent of rebased commit
+    Working copy now at: vruxwmqv bff4a4eb merge | merge
+    Parent commit      : royxmykx c84e900d b | b
+    Parent commit      : zsuskuln d57db87b a | a
+    Added 0 files, modified 0 files, removed 1 files
+    "###);
+    insta::assert_snapshot!(get_log_output(&test_env, &repo_path), @r###"
+    ◉  base
+    │ @  merge
+    ╭─┤
+    ◉ │  a
+    │ ◉  b
+    ├─╯
+    ◉
+    "###);
+
+    // Now, let's rebase onto the descendant merge
+    let (stdout, stderr) = test_env.jj_cmd_ok(&repo_path, &["op", "restore", &setup_opid]);
+    insta::assert_snapshot!(stdout, @"");
+    insta::assert_snapshot!(stderr, @r###"
+    Working copy now at: vruxwmqv b05964d1 merge | merge
+    Parent commit      : royxmykx cea87a87 b | b
+    Parent commit      : zsuskuln 2c5b7858 a | a
+    Added 1 files, modified 0 files, removed 0 files
+    "###);
+    let (stdout, stderr) = test_env.jj_cmd_ok(&repo_path, &["rebase", "-r", "base", "-d", "merge"]);
+    insta::assert_snapshot!(stdout, @"");
+    insta::assert_snapshot!(stderr, @r###"
+    Also rebased 3 descendant commits onto parent of rebased commit
+    Working copy now at: vruxwmqv 986b7a49 merge | merge
+    Parent commit      : royxmykx c07c677c b | b
+    Parent commit      : zsuskuln abc90087 a | a
+    Added 0 files, modified 0 files, removed 1 files
+    "###);
+    insta::assert_snapshot!(get_log_output(&test_env, &repo_path), @r###"
+    ◉  base
+    @    merge
+    ├─╮
     │ ◉  a
+    ◉ │  b
     ├─╯
     ◉
     "###);
+
+    // TODO(ilyagr): These will be good tests for `jj rebase --insert-after` and
+    // `--insert-before`, once those are implemented.
 }
 
 #[test]

lib/src/repo.rs

@@ -834,14 +834,32 @@ impl MutableRepo {
         )
     }
 
-    pub fn rebase_descendants(&mut self, settings: &UserSettings) -> Result<usize, TreeMergeError> {
+    fn rebase_descendants_return_rebaser<'settings, 'repo>(
+        &'repo mut self,
+        settings: &'settings UserSettings,
+    ) -> Result<Option<DescendantRebaser<'settings, 'repo>>, TreeMergeError> {
         if !self.has_rewrites() {
             // Optimization
-            return Ok(0);
+            return Ok(None);
         }
         let mut rebaser = self.create_descendant_rebaser(settings);
         rebaser.rebase_all()?;
-        Ok(rebaser.rebased().len())
+        Ok(Some(rebaser))
+    }
+
+    pub fn rebase_descendants(&mut self, settings: &UserSettings) -> Result<usize, TreeMergeError> {
+        Ok(self
+            .rebase_descendants_return_rebaser(settings)?
+            .map_or(0, |rebaser| rebaser.rebased().len()))
+    }
+
+    pub fn rebase_descendants_return_map(
+        &mut self,
+        settings: &UserSettings,
+    ) -> Result<HashMap<CommitId, CommitId>, TreeMergeError> {
+        Ok(self
+            .rebase_descendants_return_rebaser(settings)?
+            .map_or(HashMap::new(), |rebaser| rebaser.rebased().clone()))
     }
 
     pub fn set_wc_commit(