rebase: add --insert-after and --insert-before options

cli/src/commands/rebase.rs

@@ -15,6 +15,7 @@
 use std::borrow::Borrow;
 use std::collections::HashMap;
 use std::io::Write;
+use std::rc::Rc;
 use std::sync::Arc;
 
 use clap::ArgGroup;
@@ -125,6 +126,7 @@ use crate::ui::Ui;
 #[derive(clap::Args, Clone, Debug)]
 #[command(verbatim_doc_comment)]
 #[command(group(ArgGroup::new("to_rebase").args(&["branch", "source", "revision"])))]
+#[command(group(ArgGroup::new("order").args(&["destination", "insert_after", "insert_before"]).required(true)))]
 pub(crate) struct RebaseArgs {
     /// Rebase the whole branch relative to destination's ancestors (can be
     /// repeated)
@@ -157,8 +159,32 @@ pub(crate) struct RebaseArgs {
     revision: Option<RevisionArg>,
     /// The revision(s) to rebase onto (can be repeated to create a merge
     /// commit)
-    #[arg(long, short, required = true)]
+    #[arg(long, short)]
     destination: Vec<RevisionArg>,
+    /// The revision(s) to insert after (can be repeated to create a merge
+    /// commit)
+    ///
+    /// Only works with `-r`.
+    #[arg(
+        long,
+        short = 'A',
+        visible_alias = "after",
+        conflicts_with = "source",
+        conflicts_with = "branch"
+    )]
+    insert_after: Vec<RevisionArg>,
+    /// The revision(s) to insert before (can be repeated to create a merge
+    /// commit)
+    ///
+    /// Only works with `-r`.
+    #[arg(
+        long,
+        short = 'B',
+        visible_alias = "before",
+        conflicts_with = "source",
+        conflicts_with = "branch"
+    )]
+    insert_before: Vec<RevisionArg>,
 
     /// If true, when rebasing would produce an empty commit, the commit is
     /// abandoned. It will not be abandoned if it was already empty before the
@@ -193,10 +219,15 @@ Please use `jj rebase -d 'all:x|y'` instead of `jj rebase --allow-large-revsets
         simplify_ancestor_merge: false,
     };
     let mut workspace_command = command.workspace_helper(ui)?;
-    let new_parents = workspace_command
-        .resolve_some_revsets_default_single(&args.destination)?
-        .into_iter()
-        .collect_vec();
+    let new_parents = if args.destination.is_empty() {
+        vec![]
+    } else {
+        workspace_command
+            .resolve_some_revsets_default_single(&args.destination)?
+            .into_iter()
+            .collect_vec()
+    };
+
     if let Some(rev_str) = &args.revision {
         assert_eq!(
             // In principle, `-r --skip-empty` could mean to abandon the `-r`
@@ -213,13 +244,35 @@ Please use `jj rebase -d 'all:x|y'` instead of `jj rebase --allow-large-revsets
             EmptyBehaviour::Keep,
             "clap should forbid `-r --skip-empty`"
         );
-        rebase_revision(
-            ui,
-            command.settings(),
-            &mut workspace_command,
-            &new_parents,
-            rev_str,
-        )?;
+        if !args.insert_after.is_empty() {
+            let after_commits =
+                workspace_command.resolve_some_revsets_default_single(&args.insert_after)?;
+            rebase_revision_after(
+                ui,
+                command.settings(),
+                &mut workspace_command,
+                &after_commits,
+                rev_str,
+            )?;
+        } else if !args.insert_before.is_empty() {
+            let before_commits =
+                workspace_command.resolve_some_revsets_default_single(&args.insert_before)?;
+            rebase_revision_before(
+                ui,
+                command.settings(),
+                &mut workspace_command,
+                &before_commits,
+                rev_str,
+            )?;
+        } else {
+            rebase_revision(
+                ui,
+                command.settings(),
+                &mut workspace_command,
+                &new_parents,
+                rev_str,
+            )?;
+        }
     } else if !args.source.is_empty() {
         let source_commits = workspace_command.resolve_some_revsets_default_single(&args.source)?;
         rebase_descendants_transaction(
@@ -433,6 +486,85 @@ fn rebase_revision(
     }
 }
 
+fn rebase_revision_after(
+    ui: &mut Ui,
+    settings: &UserSettings,
+    workspace_command: &mut WorkspaceCommandHelper,
+    after_commits: &IndexSet<Commit>,
+    rev_str: &str,
+) -> Result<(), CommandError> {
+    let old_commit = workspace_command.resolve_single_rev(rev_str)?;
+    workspace_command.check_rewritable([&old_commit])?;
+
+    let after_commit_ids = after_commits.iter().map(|c| c.id().clone()).collect_vec();
+    let new_parents_expression = RevsetExpression::commits(after_commit_ids);
+    let new_children_expression = new_parents_expression.children();
+
+    ensure_no_commit_loop(
+        workspace_command.repo().as_ref(),
+        &new_children_expression,
+        &new_parents_expression,
+    )?;
+
+    let new_parents = after_commits;
+    let new_children: IndexSet<Commit> = new_children_expression
+        .evaluate_programmatic(workspace_command.repo().as_ref())?
+        .iter()
+        .commits(workspace_command.repo().store())
+        .try_collect()?;
+    workspace_command.check_rewritable(&new_children)?;
+
+    move_commit(
+        ui,
+        settings,
+        workspace_command,
+        new_parents,
+        &new_children,
+        old_commit,
+    )
+}
+
+fn rebase_revision_before(
+    ui: &mut Ui,
+    settings: &UserSettings,
+    workspace_command: &mut WorkspaceCommandHelper,
+    before_commits: &IndexSet<Commit>,
+    rev_str: &str,
+) -> Result<(), CommandError> {
+    let old_commit = workspace_command.resolve_single_rev(rev_str)?;
+    workspace_command.check_rewritable([&old_commit])?;
+    workspace_command.check_rewritable(before_commits.iter())?;
+
+    let before_commit_ids = before_commits.iter().map(|c| c.id().clone()).collect_vec();
+    let new_children_expression = RevsetExpression::commits(before_commit_ids);
+    let new_parents_expression = new_children_expression.parents();
+
+    ensure_no_commit_loop(
+        workspace_command.repo().as_ref(),
+        &new_children_expression,
+        &new_parents_expression,
+    )?;
+
+    let new_parents: IndexSet<Commit> = new_parents_expression
+        // Exclude parents that are ancestors of each other.
+        .minus(&new_parents_expression.parents().ancestors())
+        .clone()
+        .evaluate_programmatic(workspace_command.repo().as_ref())?
+        .iter()
+        .commits(workspace_command.repo().store())
+        .try_collect()?;
+    let new_children = before_commits;
+
+    move_commit(
+        ui,
+        settings,
+        workspace_command,
+        &new_parents,
+        new_children,
+        old_commit,
+    )
+}
+
 /// Extracts `commit` from the graph by rebasing its children on its parents.
 /// This assumes that it can be rewritten.
 fn extract_commit(
@@ -491,6 +623,155 @@ fn extract_commit(
     rebase_onto_new_parents(tx, settings, &commits_to_rebase)
 }
 
+/// Moves `commit` from its current location to a new location, given by the set
+/// of `new_parents` and `new_children_with_parents`.
+fn move_commit(
+    ui: &mut Ui,
+    settings: &UserSettings,
+    workspace_command: &mut WorkspaceCommandHelper,
+    new_parents: &IndexSet<Commit>,
+    new_children: &IndexSet<Commit>,
+    commit: Commit,
+) -> Result<(), CommandError> {
+    let (new_parents, new_children_with_parents) = move_revision_compute_parent_children(
+        workspace_command,
+        &commit,
+        new_parents,
+        new_children,
+    )?;
+
+    let mut tx = workspace_command.start_transaction();
+    // Extract `commit` from its previous location by rebasing its children
+    // onto its parents.
+    let (mut rebased_commit_ids, _) = extract_commit(&mut tx, settings, &commit)?;
+
+    // We now update `new_parents` to account for the rebase of all of
+    // `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.
+    //
+    // See comment in `rebase_revision` for the bug in this code.
+    let new_parents: Vec<_> = new_parents
+        .iter()
+        .map(|new_parent| get_possibly_rewritten_commit(&mut tx, &rebased_commit_ids, new_parent))
+        .try_collect()?;
+
+    // Finally, it's safe to rebase `commit`. We can skip rebasing if it is
+    // already a child of `new_parents`. Otherwise, at this point, it should no
+    // longer have any children; they have all been rebased and the originals
+    // have been abandoned.
+    if commit.parents() == new_parents {
+        write!(ui.stderr(), "Skipping rebase of commit ")?;
+        tx.write_commit_summary(ui.stderr_formatter().as_mut(), &commit)?;
+        writeln!(ui.stderr())?;
+    } else {
+        let new_commit = rebase_commit(settings, tx.mut_repo(), &commit, &new_parents)?;
+        rebased_commit_ids.insert(commit.id().clone(), new_commit.id().clone());
+    };
+
+    // Now, rebase all the new children onto the newly rebased commit.
+    // TODO: Make this no-op if the new children already have the required parents.
+    let new_children_with_parents: Vec<_> = new_children_with_parents
+        .iter()
+        .map(|(commit, parents)| {
+            // Both the original child commit and its parents could have been rewritten.
+            let commit = get_possibly_rewritten_commit(&mut tx, &rebased_commit_ids, commit)?;
+            let new_parents: Vec<_> = parents
+                .iter()
+                .map(|parent| get_possibly_rewritten_commit(&mut tx, &rebased_commit_ids, parent))
+                .try_collect()?;
+            Ok::<_, BackendError>((commit, new_parents))
+        })
+        .try_collect()?;
+
+    let (new_rebased_commit_ids, _) =
+        rebase_onto_new_parents(&mut tx, settings, &new_children_with_parents)?;
+
+    let num_rebased_commits = rebased_commit_ids
+        .iter()
+        .filter(|(_, new_commit_id)| !new_rebased_commit_ids.contains_key(new_commit_id))
+        .count()
+        + new_rebased_commit_ids.len();
+
+    if num_rebased_commits > 0 {
+        writeln!(ui.stderr(), "Rebased {num_rebased_commits} commits")?;
+    }
+    if tx.mut_repo().has_changes() {
+        tx.finish(ui, format!("move commit {}", commit.id().hex()))
+    } else {
+        Ok(()) // Do not print "Nothing changed."
+    }
+}
+
+type NewParentsAndChildren = (Vec<Commit>, Vec<(Commit, Vec<Commit>)>);
+
+fn move_revision_compute_parent_children(
+    workspace_command: &mut WorkspaceCommandHelper,
+    old_commit: &Commit,
+    new_parents_set: &IndexSet<Commit>,
+    new_children_set: &IndexSet<Commit>,
+) -> Result<NewParentsAndChildren, CommandError> {
+    // If the new parents include the old commit, replace it with the old commit's
+    // parents.
+    // e.g. `jj rebase -r A --before A`
+    let new_parents: Vec<_> = new_parents_set
+        .iter()
+        .flat_map(|c| {
+            if c == old_commit {
+                old_commit.parents().clone()
+            } else {
+                [c.clone()].to_vec()
+            }
+        })
+        .collect();
+
+    let old_commit_expression = RevsetExpression::commit(old_commit.id().clone());
+    let old_commit_children: Vec<Commit> = old_commit_expression
+        .children()
+        .evaluate_programmatic(workspace_command.repo().as_ref())?
+        .iter()
+        .commits(workspace_command.repo().store())
+        .try_collect()?;
+
+    // If the new children include the old commit, replace it with the old commit's
+    // children.
+    // e.g. `jj rebase -r A --after A`
+    let new_children: Vec<_> = new_children_set
+        .iter()
+        .flat_map(|c| {
+            if c == old_commit {
+                old_commit_children.clone()
+            } else {
+                [c.clone()].to_vec()
+            }
+        })
+        .collect();
+
+    let new_children_to_rebase: Vec<_> = new_children
+        .into_iter()
+        .map(|child_commit| {
+            // Exclude any of the new parents, since we are "inserting" the new commit
+            // in between the new parents and the new children.
+            let mut new_child_parents: Vec<_> = child_commit
+                .parent_ids()
+                .iter()
+                .filter(|&id| !new_parents_set.iter().any(|commit| commit.id() == id))
+                .map(|id| workspace_command.repo().store().get_commit(id))
+                .try_collect()?;
+
+            // Add the commit as a parent of the new child commit if it is not already
+            // a parent.
+            if !new_child_parents.iter().any(|c| c == old_commit) {
+                new_child_parents.push(old_commit.clone());
+            }
+
+            Ok::<_, BackendError>((child_commit, new_child_parents))
+        })
+        .try_collect()?;
+
+    Ok((new_parents, new_children_to_rebase))
+}
+
 /// Rebases each commit in `commits_to_rebase` onto its new parents, and calls
 /// `rebase_descendants` after.
 /// This assumes that each commit can be rewritten.
@@ -534,6 +815,28 @@ fn get_possibly_rewritten_commit(
         })
 }
 
+/// Ensure that there is no possible cycle between the potential children and
+/// parents of a rebased commit.
+fn ensure_no_commit_loop(
+    repo: &ReadonlyRepo,
+    children_expression: &Rc<RevsetExpression>,
+    parents_expression: &Rc<RevsetExpression>,
+) -> Result<(), CommandError> {
+    if let Some(commit_id) = children_expression
+        .dag_range_to(parents_expression)
+        .evaluate_programmatic(repo)?
+        .iter()
+        .next()
+    {
+        return Err(user_error(format!(
+            "Refusing to create a loop: commit {} would be both an ancestor and a descendant of \
+             the rebased commit",
+            short_commit_hash(&commit_id),
+        )));
+    }
+    Ok(())
+}
+
 fn check_rebase_destinations(
     repo: &Arc<ReadonlyRepo>,
     new_parents: &[Commit],