cli: migrate custom command augmentation and parse_args() to CliRunner

CommandHelper is passed by move because custom-backend/main.rs mutates it.
This will be addressed by the next commit.

examples/custom-backend/main.rs

@@ -15,11 +15,10 @@
 use std::io::Read;
 use std::path::Path;
 
-use clap::{FromArgMatches, Subcommand};
+use clap::{ArgMatches, FromArgMatches};
 use git2::Repository;
-use jujutsu::cli_util::{parse_args, CliRunner, CommandError, TracingSubscription};
-use jujutsu::commands::{default_app, run_command};
-use jujutsu::config::read_config;
+use jujutsu::cli_util::{CliRunner, CommandError, CommandHelper};
+use jujutsu::commands::run_command;
 use jujutsu::ui::Ui;
 use jujutsu_lib::backend::{
     Backend, BackendResult, Commit, CommitId, Conflict, ConflictId, FileId, SymlinkId, Tree, TreeId,
@@ -35,11 +34,11 @@ enum CustomCommands {
     InitJit,
 }
 
-fn run(ui: &mut Ui, tracing_subscription: &TracingSubscription) -> Result<(), CommandError> {
-    ui.reset(read_config()?);
-    let app = CustomCommands::augment_subcommands(default_app());
-    let (mut command_helper, matches) =
-        parse_args(ui, app, tracing_subscription, std::env::args_os())?;
+fn run(
+    ui: &mut Ui,
+    mut command_helper: CommandHelper,
+    matches: &ArgMatches,
+) -> Result<(), CommandError> {
     let mut store_factories = StoreFactories::default();
     // Register the backend so it can be loaded when the repo is loaded. The name
     // must match `Backend::name()`.
@@ -48,7 +47,7 @@ fn run(ui: &mut Ui, tracing_subscription: &TracingSubscription) -> Result<(), Co
         Box::new(|store_path| Box::new(JitBackend::load(store_path))),
     );
     command_helper.set_store_factories(store_factories);
-    match CustomCommands::from_arg_matches(&matches) {
+    match CustomCommands::from_arg_matches(matches) {
         // Handle our custom command
         Ok(CustomCommands::InitJit) => {
             let wc_path = ui.cwd();
@@ -59,12 +58,15 @@ fn run(ui: &mut Ui, tracing_subscription: &TracingSubscription) -> Result<(), Co
             Ok(())
         }
         // Handle default commands
-        Err(_) => run_command(ui, &command_helper, &matches),
+        Err(_) => run_command(ui, &command_helper, matches),
     }
 }
 
 fn main() {
-    CliRunner::init().set_dispatch_fn(run).run_and_exit();
+    CliRunner::init()
+        .add_subcommand::<CustomCommands>()
+        .set_dispatch_fn(run)
+        .run_and_exit();
 }
 
 /// A commit backend that's extremely similar to the Git backend

examples/custom-command/main.rs

@@ -12,12 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use clap::{FromArgMatches, Subcommand};
-use jujutsu::cli_util::{
-    parse_args, short_commit_description, CliRunner, CommandError, TracingSubscription,
-};
-use jujutsu::commands::{default_app, run_command};
-use jujutsu::config::read_config;
+use clap::{ArgMatches, FromArgMatches};
+use jujutsu::cli_util::{short_commit_description, CliRunner, CommandError, CommandHelper};
+use jujutsu::commands::run_command;
 use jujutsu::ui::Ui;
 
 #[derive(clap::Parser, Clone, Debug)]
@@ -33,11 +30,12 @@ struct FrobnicateArgs {
     revision: String,
 }
 
-fn run(ui: &mut Ui, tracing_subscription: &TracingSubscription) -> Result<(), CommandError> {
-    ui.reset(read_config()?);
-    let app = CustomCommands::augment_subcommands(default_app());
-    let (command_helper, matches) = parse_args(ui, app, tracing_subscription, std::env::args_os())?;
-    match CustomCommands::from_arg_matches(&matches) {
+fn run(
+    ui: &mut Ui,
+    command_helper: CommandHelper,
+    matches: &ArgMatches,
+) -> Result<(), CommandError> {
+    match CustomCommands::from_arg_matches(matches) {
         // Handle our custom command
         Ok(CustomCommands::Frobnicate(args)) => {
             let mut workspace_command = command_helper.workspace_helper(ui)?;
@@ -57,10 +55,13 @@ fn run(ui: &mut Ui, tracing_subscription: &TracingSubscription) -> Result<(), Co
             Ok(())
         }
         // Handle default commands
-        Err(_) => run_command(ui, &command_helper, &matches),
+        Err(_) => run_command(ui, &command_helper, matches),
     }
 }
 
 fn main() {
-    CliRunner::init().set_dispatch_fn(run).run_and_exit();
+    CliRunner::init()
+        .add_subcommand::<CustomCommands>()
+        .set_dispatch_fn(run)
+        .run_and_exit();
 }

src/cli_util.rs

@@ -1687,6 +1687,7 @@ pub fn handle_command_result(ui: &mut Ui, result: Result<(), CommandError>) -> i
 #[must_use]
 pub struct CliRunner<F> {
     tracing_subscription: TracingSubscription,
+    app: clap::Command,
     dispatch_fn: F,
 }
 
@@ -1698,17 +1699,32 @@ impl CliRunner<()> {
         crate::cleanup_guard::init();
         CliRunner {
             tracing_subscription,
+            app: crate::commands::default_app(),
             dispatch_fn: (),
         }
     }
 
+    /// Registers new subcommands in addition to the default ones.
+    // TODO: maybe take dispatch_fn for the subcommands?
+    pub fn add_subcommand<C>(self) -> Self
+    where
+        C: clap::Subcommand,
+    {
+        CliRunner {
+            tracing_subscription: self.tracing_subscription,
+            app: C::augment_subcommands(self.app),
+            dispatch_fn: self.dispatch_fn,
+        }
+    }
+
     // TODO: use crate::commands::run_command() by default
     pub fn set_dispatch_fn<F>(self, dispatch_fn: F) -> CliRunner<F>
     where
-        F: FnOnce(&mut Ui, &TracingSubscription) -> Result<(), CommandError>,
+        F: FnOnce(&mut Ui, CommandHelper, &ArgMatches) -> Result<(), CommandError>,
     {
         CliRunner {
             tracing_subscription: self.tracing_subscription,
+            app: self.app,
             dispatch_fn,
         }
     }
@@ -1716,10 +1732,18 @@ impl CliRunner<()> {
 
 impl<F> CliRunner<F>
 where
-    F: FnOnce(&mut Ui, &TracingSubscription) -> Result<(), CommandError>,
+    F: FnOnce(&mut Ui, CommandHelper, &ArgMatches) -> Result<(), CommandError>,
 {
     pub fn run(self, ui: &mut Ui) -> Result<(), CommandError> {
-        (self.dispatch_fn)(ui, &self.tracing_subscription)
+        ui.reset(crate::config::read_config()?);
+        let (command_helper, matches) = parse_args(
+            ui,
+            self.app,
+            &self.tracing_subscription,
+            std::env::args_os(),
+        )?;
+        // TODO: pass CommandHelper by reference
+        (self.dispatch_fn)(ui, command_helper, &matches)
     }
 
     pub fn run_and_exit(self) -> ! {

src/main.rs

@@ -12,18 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use jujutsu::cli_util::{parse_args, CliRunner, CommandError, TracingSubscription};
-use jujutsu::commands::{default_app, run_command};
-use jujutsu::config::read_config;
-use jujutsu::ui::Ui;
-
-fn run(ui: &mut Ui, tracing_subscription: &TracingSubscription) -> Result<(), CommandError> {
-    ui.reset(read_config()?);
-    let app = default_app();
-    let (command_helper, matches) = parse_args(ui, app, tracing_subscription, std::env::args_os())?;
-    run_command(ui, &command_helper, &matches)
-}
+use jujutsu::cli_util::CliRunner;
+use jujutsu::commands::run_command;
 
 fn main() {
-    CliRunner::init().set_dispatch_fn(run).run_and_exit();
+    CliRunner::init()
+        .set_dispatch_fn(|ui, command_helper, matches| run_command(ui, &command_helper, matches))
+        .run_and_exit();
 }