From d7d156e7e6c9fa4dc9a545ca898cbb28ae9f1c98 Mon Sep 17 00:00:00 2001
From: Evan Mesterhazy <emesterhazy@google.com>
Date: Mon, 18 Mar 2024 23:35:10 -0400
Subject: [PATCH] Implement advance-branches for jj new

---
 cli/src/commands/new.rs            | 31 +++++++++++++++++++++++++++---
 cli/tests/test_advance_branches.rs |  9 +++++++++
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/cli/src/commands/new.rs b/cli/src/commands/new.rs
index a94d80769e2..98ebd3479de 100644
--- a/cli/src/commands/new.rs
+++ b/cli/src/commands/new.rs
@@ -22,7 +22,9 @@ use jj_lib::revset::{RevsetExpression, RevsetIteratorExt};
 use jj_lib::rewrite::{merge_commit_trees, rebase_commit};
 use tracing::instrument;
 
-use crate::cli_util::{self, short_commit_hash, CommandHelper, RevisionArg};
+use crate::cli_util::{
+    self, short_commit_hash, AdvanceableBranch, CommandHelper, RevisionArg, WorkspaceCommandHelper,
+};
 use crate::command_error::{user_error, CommandError};
 use crate::description_util::join_message_paragraphs;
 use crate::ui::Ui;
@@ -102,6 +104,8 @@ Please use `jj new 'all:x|y'` instead of `jj new --allow-large-revsets x y`.",
         .into_iter()
         .collect_vec();
     let target_ids = target_commits.iter().map(|c| c.id().clone()).collect_vec();
+    let advanceable_branches =
+        get_advanceable_branches(args, &workspace_command, &target_commits);
     let mut tx = workspace_command.start_transaction();
     let mut num_rebased;
     let new_commit;
@@ -145,11 +149,11 @@ Please use `jj new 'all:x|y'` instead of `jj new --allow-large-revsets x y`.",
             .set_description(join_message_paragraphs(&args.message_paragraphs))
             .write()?;
         num_rebased = target_ids.len();
-        for child_commit in target_commits {
+        for child_commit in &target_commits {
             rebase_commit(
                 command.settings(),
                 tx.mut_repo(),
-                &child_commit,
+                child_commit,
                 &[new_commit.clone()],
             )?;
         }
@@ -206,6 +210,27 @@ Please use `jj new 'all:x|y'` instead of `jj new --allow-large-revsets x y`.",
     if num_rebased > 0 {
         writeln!(ui.stderr(), "Rebased {num_rebased} descendant commits")?;
     }
+
+    if !advanceable_branches.is_empty() {
+        tx.advance_branches(advanceable_branches, target_commits[0].id());
+    }
+
     tx.finish(ui, "new empty commit")?;
     Ok(())
 }
+
+// Branches are only advanced if `jj new` has a single target commit and the
+// new commit is not being inserted before or after the target.
+fn get_advanceable_branches(
+    args: &NewArgs,
+    ws: &WorkspaceCommandHelper,
+    target_commits: &[Commit],
+) -> Vec<AdvanceableBranch> {
+    let should_advance_branches = target_commits.len() == 1 &&
+        !args.insert_before && !args.insert_after;
+    if should_advance_branches {
+        ws.get_advanceable_branches(&**ws.repo(), target_commits[0].parent_ids())
+    } else {
+        Vec::new()
+    }
+}
diff --git a/cli/tests/test_advance_branches.rs b/cli/tests/test_advance_branches.rs
index a95eb72cc0c..936393ef4e4 100644
--- a/cli/tests/test_advance_branches.rs
+++ b/cli/tests/test_advance_branches.rs
@@ -52,6 +52,11 @@ fn commit_cmd(env: &TestEnvironment, workspace_path: &Path, commit_message: &str
     env.jj_cmd_ok(workspace_path, &["commit", "-m", commit_message]);
 }
 
+fn describe_new_cmd(env: &TestEnvironment, workspace_path: &Path, commit_message: &str) {
+    env.jj_cmd_ok(workspace_path, &["describe", "-m", commit_message]);
+    env.jj_cmd_ok(workspace_path, &["new"]);
+}
+
 macro_rules! parameterized_tests{
     ($($test_name:ident: $case_fn:ident($commit_fn:ident),)*) => {
     $(
@@ -68,6 +73,10 @@ parameterized_tests! {
     test_commit_advance_branches_at_minus: case_advance_branches_at_minus(commit_cmd),
     test_commit_advance_branches_overrides: case_advance_branches_overrides(commit_cmd),
     test_commit_advance_branches_ambiguity: case_advance_branches_ambiguity(commit_cmd),
+    test_new_advance_branches_enabled: case_advance_branches_enabled(describe_new_cmd),
+    test_new_advance_branches_at_minus: case_advance_branches_at_minus(describe_new_cmd),
+    test_new_advance_branches_overrides: case_advance_branches_overrides(describe_new_cmd),
+    test_new_advance_branches_ambiguity: case_advance_branches_ambiguity(describe_new_cmd),
 }
 
 // Check that enabling and disabling advance-branches works as expected.