From 110ae76365d1b581c5b9a28523a713e052c05f4d Mon Sep 17 00:00:00 2001
From: Alec Snyder <linuxbash8@gmail.com>
Date: Wed, 15 May 2024 23:24:00 +0300
Subject: [PATCH] Add a branch option to 'jj new'

---
 cli/src/commands/new.rs          | 17 +++++++++++++++
 cli/tests/cli-reference@.md.snap |  1 +
 cli/tests/test_new_command.rs    | 36 ++++++++++++++++++++++++++++++++
 3 files changed, 54 insertions(+)

diff --git a/cli/src/commands/new.rs b/cli/src/commands/new.rs
index 9655b7bb57..04e8b2e2f2 100644
--- a/cli/src/commands/new.rs
+++ b/cli/src/commands/new.rs
@@ -17,6 +17,7 @@ use std::io::Write;
 use clap::ArgGroup;
 use itertools::Itertools;
 use jj_lib::commit::{Commit, CommitIteratorExt};
+use jj_lib::op_store::RefTarget;
 use jj_lib::repo::Repo;
 use jj_lib::revset::{RevsetExpression, RevsetIteratorExt};
 use jj_lib::rewrite::{merge_commit_trees, rebase_commit};
@@ -50,6 +51,9 @@ pub(crate) struct NewArgs {
     /// The change description to use
     #[arg(long = "message", short, value_name = "MESSAGE")]
     message_paragraphs: Vec<String>,
+    /// Create a new branch or set an existing branch to the new commit
+    #[arg(long = "branch", short, value_name = "BRANCH")]
+    branch_name: Option<String>,
     /// Deprecated. Please prefix the revset with `all:` instead.
     #[arg(long, short = 'L', hide = true)]
     allow_large_revsets: bool,
@@ -215,6 +219,19 @@ Please use `jj new 'all:x|y'` instead of `jj new --allow-large-revsets x y`.",
         tx.advance_branches(advanceable_branches, &target);
     }
 
+    if let Some(branch_name) = &args.branch_name {
+        let view = tx.repo().view();
+        if view.get_local_branch(branch_name).is_present() {
+            writeln!(
+                ui.warning_default(),
+                "setting existing branch: {}",
+                branch_name,
+            )?;
+        }
+        tx.mut_repo()
+            .set_local_branch_target(branch_name, RefTarget::normal(new_commit.id().clone()));
+    }
+
     tx.finish(ui, "new empty commit")?;
     Ok(())
 }
diff --git a/cli/tests/cli-reference@.md.snap b/cli/tests/cli-reference@.md.snap
index c74cbbfa50..103edb7dc9 100644
--- a/cli/tests/cli-reference@.md.snap
+++ b/cli/tests/cli-reference@.md.snap
@@ -1130,6 +1130,7 @@ For more information, see https://github.com/martinvonz/jj/blob/main/docs/workin
 
 * `-r` — Ignored (but lets you pass `-r` for consistency with other commands)
 * `-m`, `--message <MESSAGE>` — The change description to use
+* `-b`, `--branch <BRANCH>` — Create a new branch or set an existing branch to the new commit
 * `-L`, `--allow-large-revsets` — Deprecated. Please prefix the revset with `all:` instead
 
   Possible values: `true`, `false`
diff --git a/cli/tests/test_new_command.rs b/cli/tests/test_new_command.rs
index 4eace54921..59e4ab7c9e 100644
--- a/cli/tests/test_new_command.rs
+++ b/cli/tests/test_new_command.rs
@@ -529,6 +529,37 @@ fn test_new_error_revision_does_not_exist() {
     "###);
 }
 
+#[test]
+fn test_new_with_new_branch() {
+    let test_env = TestEnvironment::default();
+    test_env.jj_cmd_ok(test_env.env_root(), &["git", "init", "repo"]);
+    let repo_path = test_env.env_root().join("repo");
+    
+    test_env.jj_cmd_ok(&repo_path, &["new", "-b", "new_branch"]);
+
+    insta::assert_snapshot!(get_log_output_branches(&test_env, &repo_path), @r###"
+    @  65b6b74e08973b88d38404430f119c8c79465250 new_branch
+    ◉  230dd059e1b059aefc0da06a2e5a7dbf22362f22
+    ◉  0000000000000000000000000000000000000000
+    "###);
+}
+
+#[test]
+fn test_new_with_existing_branch() {
+    let test_env = TestEnvironment::default();
+    test_env.jj_cmd_ok(test_env.env_root(), &["git", "init", "repo"]);
+    let repo_path = test_env.env_root().join("repo");
+
+    test_env.jj_cmd_ok(&repo_path, &["branch", "c", "new_branch"]);
+    test_env.jj_cmd_ok(&repo_path, &["new", "-b", "new_branch"]);
+
+    insta::assert_snapshot!(get_log_output_branches(&test_env, &repo_path), @r###"
+    @  4db490c88528133d579540b6900b8098f0c17701 new_branch
+    ◉  230dd059e1b059aefc0da06a2e5a7dbf22362f22
+    ◉  0000000000000000000000000000000000000000
+    "###);
+}
+
 fn setup_before_insertion(test_env: &TestEnvironment, repo_path: &Path) {
     test_env.jj_cmd_ok(repo_path, &["branch", "create", "A"]);
     test_env.jj_cmd_ok(repo_path, &["commit", "-m", "A"]);
@@ -550,6 +581,11 @@ fn get_log_output(test_env: &TestEnvironment, repo_path: &Path) -> String {
     test_env.jj_cmd_success(repo_path, &["log", "-T", template])
 }
 
+fn get_log_output_branches(test_env: &TestEnvironment, repo_path: &Path) -> String {
+    let template = r#"commit_id ++ " " ++ branches"#;
+    test_env.jj_cmd_success(repo_path, &["log", "-T", template])
+}
+
 fn get_short_log_output(test_env: &TestEnvironment, repo_path: &Path) -> String {
     let template = r#"if(description, description, "root")"#;
     test_env.jj_cmd_success(repo_path, &["log", "-T", template])