From 45ebaacbed2849c5956065f617132ede84f6cbd6 Mon Sep 17 00:00:00 2001
From: Evan Simmons <esims89@gmail.com>
Date: Thu, 7 Mar 2024 20:41:28 -0800
Subject: [PATCH] add branch batch create

---
 cli/src/commands/branch.rs | 104 +++++++++++++++++++++++++++++++++++--
 1 file changed, 99 insertions(+), 5 deletions(-)

diff --git a/cli/src/commands/branch.rs b/cli/src/commands/branch.rs
index 9f936ffedaa..678816a21e9 100644
--- a/cli/src/commands/branch.rs
+++ b/cli/src/commands/branch.rs
@@ -23,15 +23,16 @@ use jj_lib::git;
 use jj_lib::object_id::ObjectId;
 use jj_lib::op_store::{RefTarget, RemoteRef};
 use jj_lib::repo::Repo;
-use jj_lib::revset::{self, RevsetExpression};
+use jj_lib::revset::{self, RevsetExpression, RevsetIteratorExt};
 use jj_lib::str_util::StringPattern;
 use jj_lib::view::View;
 
 use crate::cli_util::{
-    parse_string_pattern, CommandHelper, RemoteBranchName, RemoteBranchNamePattern, RevisionArg,
+    self, parse_string_pattern, CommandHelper, RemoteBranchName, RemoteBranchNamePattern,
+    RevisionArg,
 };
 use crate::command_error::{user_error, user_error_with_hint, CommandError};
-use crate::formatter::Formatter;
+use crate::formatter::{self, Formatter};
 use crate::ui::Ui;
 
 /// Manage branches.
@@ -40,6 +41,8 @@ use crate::ui::Ui;
 /// https://github.com/martinvonz/jj/blob/main/docs/branches.md.
 #[derive(clap::Subcommand, Clone, Debug)]
 pub enum BranchCommand {
+    #[command(visible_alias("b"))]
+    BatchCreate(BranchBatchCreateArgs),
     #[command(visible_alias("c"))]
     Create(BranchCreateArgs),
     #[command(visible_alias("d"))]
@@ -57,6 +60,20 @@ pub enum BranchCommand {
     Untrack(BranchUntrackArgs),
 }
 
+/// Create a new branch.
+#[derive(clap::Args, Clone, Debug)]
+pub struct BranchBatchCreateArgs {
+    /// The branch's target revision.
+    #[arg(long, short)]
+    revisions: Vec<RevisionArg>,
+
+    /// Create a branch with the given template
+    ///
+    /// For the syntax, see https://github.com/martinvonz/jj/blob/main/docs/templates.md
+    #[arg(long, short = 'T')]
+    template: Option<String>,
+}
+
 /// Create a new branch.
 #[derive(clap::Args, Clone, Debug)]
 pub struct BranchCreateArgs {
@@ -221,6 +238,7 @@ pub fn cmd_branch(
     subcommand: &BranchCommand,
 ) -> Result<(), CommandError> {
     match subcommand {
+        BranchCommand::BatchCreate(sub_args) => cmd_branch_batch_create(ui, command, sub_args),
         BranchCommand::Create(sub_args) => cmd_branch_create(ui, command, sub_args),
         BranchCommand::Rename(sub_args) => cmd_branch_rename(ui, command, sub_args),
         BranchCommand::Set(sub_args) => cmd_branch_set(ui, command, sub_args),
@@ -232,6 +250,80 @@ pub fn cmd_branch(
     }
 }
 
+fn cmd_branch_batch_create(
+    ui: &mut Ui,
+    command: &CommandHelper,
+    args: &BranchBatchCreateArgs,
+) -> Result<(), CommandError> {
+    let mut workspace_command = command.workspace_helper(ui)?;
+    let view = workspace_command.repo().view();
+
+    let revset_expression = {
+        let expression = if args.revisions.is_empty() {
+            workspace_command.parse_revset(&command.settings().default_revset())?
+        } else {
+            let expressions: Vec<_> = args
+                .revisions
+                .iter()
+                .map(|revision_str| workspace_command.parse_revset(revision_str))
+                .try_collect()?;
+            RevsetExpression::union_all(&expressions)
+        };
+
+        revset::optimize(expression)
+    };
+    let mut branches = vec![];
+    {
+        let repo = workspace_command.repo();
+        let store = repo.store();
+        let revset = workspace_command.evaluate_revset(revset_expression)?;
+        let iter: Box<dyn Iterator<Item = CommitId>> = Box::new(revset.iter());
+        let template =
+            workspace_command.parse_commit_template(args.template.as_deref().unwrap())?;
+
+        for commit_or_error in iter.commits(store) {
+            let commit = commit_or_error?;
+            let mut fmtr = formatter::FormatRecorder::new();
+            template.format(&commit, &mut fmtr)?;
+            let data = fmtr.data();
+            let branch_name = std::str::from_utf8(data).unwrap();
+            branches.push((branch_name.to_string(), commit));
+        }
+    }
+
+    if let Some((branch_name, _)) = branches
+        .iter()
+        .find(|(name, _)| view.get_local_branch(name).is_present())
+    {
+        return Err(user_error_with_hint(
+            format!("Branch already exists: {branch_name}"),
+            "Use `jj branch set` to update it.",
+        ));
+    }
+
+    if branches.len() > 1 {
+        writeln!(
+            ui.warning(),
+            "warning: Creating multiple branches: {}",
+            branches.iter().map(|(n, _)| n).join(", "),
+        )?;
+    }
+
+    let mut tx = workspace_command.start_transaction();
+    for (branch_name, target_commit) in &branches {
+        tx.mut_repo()
+            .set_local_branch_target(branch_name, RefTarget::normal(target_commit.id().clone()));
+    }
+    tx.finish(
+        ui,
+        format!(
+            "create batch of branches {}",
+            branches.iter().map(|x| &x.0).join(",")
+        ),
+    )?;
+    Ok(())
+}
+
 fn cmd_branch_create(
     ui: &mut Ui,
     command: &CommandHelper,
@@ -241,7 +333,9 @@ fn cmd_branch_create(
     let target_commit =
         workspace_command.resolve_single_rev(args.revision.as_deref().unwrap_or("@"))?;
     let view = workspace_command.repo().view();
-    let branch_names = &args.names;
+
+    let mut branch_names = &args.names;
+
     if let Some(branch_name) = branch_names
         .iter()
         .find(|&name| view.get_local_branch(name).is_present())
@@ -269,7 +363,7 @@ fn cmd_branch_create(
         ui,
         format!(
             "create {} pointing to commit {}",
-            make_branch_term(branch_names),
+            make_branch_term(&branch_names),
             target_commit.id().hex()
         ),
     )?;