rhg: refactor to pass argv down, instead of caling args_os()

This refactoring makes it easy to patch some command-line preprocessing into rhg.
We use this to support using rhg as a shebang interpreter, for example.

Differential Revision: https://phab.mercurial-scm.org/D12543

rust/rhg/src/blackbox.rs

@@ -5,6 +5,7 b' use format_bytes::format_bytes;'
 use hg::errors::HgError;
 use hg::repo::Repo;
 use hg::utils::{files::get_bytes_from_os_str, shell_quote};
+use std::ffi::OsString;
 
 const ONE_MEBIBYTE: u64 = 1 << 20;
 
@@ -83,14 +84,21 b" impl<'a> Blackbox<'a> {"
         })
     }
 
-    pub fn log_command_start(&self) {
+    pub fn log_command_start<'arg>(
+        &self,
+        argv: impl Iterator<Item = &'arg OsString>,
+    ) {
         if let Some(configured) = &self.configured {
-            let message = format_bytes!(b"(rust) {}", format_cli_args());
+            let message = format_bytes!(b"(rust) {}", format_cli_args(argv));
             configured.log(&self.process_start_time.calendar_based, &message);
         }
     }
 
-    pub fn log_command_end(&self, exit_code: i32) {
+    pub fn log_command_end<'arg>(
+        &self,
+        argv: impl Iterator<Item = &'arg OsString>,
+        exit_code: i32,
+    ) {
         if let Some(configured) = &self.configured {
             let now = chrono::Local::now();
             let duration = self
@@ -100,7 +108,7 b" impl<'a> Blackbox<'a> {"
                 .as_secs_f64();
             let message = format_bytes!(
                 b"(rust) {} exited {} after {} seconds",
-                format_cli_args(),
+                format_cli_args(argv),
                 exit_code,
                 format_bytes::Utf8(format_args!("{:.03}", duration))
             );
@@ -147,8 +155,9 b" impl ConfiguredBlackbox<'_> {"
     }
 }
 
-fn format_cli_args() -> Vec<u8> {
-    let mut args = std::env::args_os();
+fn format_cli_args<'a>(
+    mut args: impl Iterator<Item = &'a OsString>,
+) -> Vec<u8> {
     let _ = args.next(); // Skip the first (or zeroth) arg, the name of the `rhg` executable
     let mut args = args.map(|arg| shell_quote(&get_bytes_from_os_str(arg)));
     let mut formatted = Vec::new();

rust/rhg/src/main.rs

@@ -25,6 +25,7 b' pub mod utils {'
 }
 
 fn main_with_result(
+    argv: Vec<OsString>,
     process_start_time: &blackbox::ProcessStartTime,
     ui: &ui::Ui,
     repo: Result<&Repo, &NoRepoInCwdError>,
@@ -78,7 +79,7 b' fn main_with_result('
         .version("0.0.1");
     let app = add_subcommand_args(app);
 
-    let matches = app.clone().get_matches_safe()?;
+    let matches = app.clone().get_matches_from_safe(argv.iter())?;
 
     let (subcommand_name, subcommand_matches) = matches.subcommand();
 
@@ -123,23 +124,26 b' fn main_with_result('
     if config.is_extension_enabled(b"blackbox") {
         let blackbox =
             blackbox::Blackbox::new(&invocation, process_start_time)?;
-        blackbox.log_command_start();
+        blackbox.log_command_start(argv.iter());
         let result = run(&invocation);
-        blackbox.log_command_end(exit_code(
-            &result,
-            // TODO: show a warning or combine with original error if
-            // `get_bool` returns an error
-            config
-                .get_bool(b"ui", b"detailed-exit-code")
-                .unwrap_or(false),
-        ));
+        blackbox.log_command_end(
+            argv.iter(),
+            exit_code(
+                &result,
+                // TODO: show a warning or combine with original error if
+                // `get_bool` returns an error
+                config
+                    .get_bool(b"ui", b"detailed-exit-code")
+                    .unwrap_or(false),
+            ),
+        );
         result
     } else {
         run(&invocation)
     }
 }
 
-fn main() {
+fn rhg_main(argv: Vec<OsString>) -> ! {
     // Run this first, before we find out if the blackbox extension is even
     // enabled, in order to include everything in-between in the duration
     // measurements. Reading config files can be slow if they’re on NFS.
@@ -147,7 +151,7 b' fn main() {'
 
     env_logger::init();
 
-    let early_args = EarlyArgs::parse(std::env::args_os());
+    let early_args = EarlyArgs::parse(&argv);
 
     let initial_current_dir = early_args.cwd.map(|cwd| {
         let cwd = get_path_from_bytes(&cwd);
@@ -158,6 +162,7 b' fn main() {'
             })
             .unwrap_or_else(|error| {
                 exit(
+                    &argv,
                     &None,
                     &Ui::new_infallible(&Config::empty()),
                     OnUnsupported::Abort,
@@ -179,6 +184,7 b' fn main() {'
             let on_unsupported = OnUnsupported::Abort;
 
             exit(
+                &argv,
                 &initial_current_dir,
                 &Ui::new_infallible(&Config::empty()),
                 on_unsupported,
@@ -191,6 +197,7 b' fn main() {'
         .load_cli_args(early_args.config, early_args.color)
         .unwrap_or_else(|error| {
             exit(
+                &argv,
                 &initial_current_dir,
                 &Ui::new_infallible(&non_repo_config),
                 OnUnsupported::from_config(&non_repo_config),
@@ -209,6 +216,7 b' fn main() {'
         }
         if SCHEME_RE.is_match(&repo_path_bytes) {
             exit(
+                &argv,
                 &initial_current_dir,
                 &Ui::new_infallible(&non_repo_config),
                 OnUnsupported::from_config(&non_repo_config),
@@ -299,6 +307,7 b' fn main() {'
             Err(NoRepoInCwdError { cwd: at })
         }
         Err(error) => exit(
+            &argv,
             &initial_current_dir,
             &Ui::new_infallible(&non_repo_config),
             OnUnsupported::from_config(&non_repo_config),
@@ -318,6 +327,7 b' fn main() {'
     };
     let ui = Ui::new(&config).unwrap_or_else(|error| {
         exit(
+            &argv,
             &initial_current_dir,
             &Ui::new_infallible(&config),
             OnUnsupported::from_config(&config),
@@ -330,12 +340,14 b' fn main() {'
     let on_unsupported = OnUnsupported::from_config(config);
 
     let result = main_with_result(
+        argv.iter().map(|s| s.to_owned()).collect(),
         &process_start_time,
         &ui,
         repo_result.as_ref(),
         config,
     );
     exit(
+        &argv,
         &initial_current_dir,
         &ui,
         on_unsupported,
@@ -348,6 +360,10 b' fn main() {'
     )
 }
 
+fn main() -> ! {
+    rhg_main(std::env::args_os().collect())
+}
+
 fn exit_code(
     result: &Result<(), CommandError>,
     use_detailed_exit_code: bool,
@@ -374,7 +390,8 b' fn exit_code('
     }
 }
 
-fn exit(
+fn exit<'a>(
+    original_args: &'a [OsString],
     initial_current_dir: &Option<PathBuf>,
     ui: &Ui,
     mut on_unsupported: OnUnsupported,
@@ -386,7 +403,7 b' fn exit('
         Err(CommandError::UnsupportedFeature { message }),
     ) = (&on_unsupported, &result)
     {
-        let mut args = std::env::args_os();
+        let mut args = original_args.iter();
         let executable = match executable {
             None => {
                 exit_no_fallback(
@@ -546,7 +563,7 b' struct EarlyArgs {'
 }
 
 impl EarlyArgs {
-    fn parse(args: impl IntoIterator<Item = OsString>) -> Self {
+    fn parse<'a>(args: impl IntoIterator<Item = &'a OsString>) -> Self {
         let mut args = args.into_iter().map(get_bytes_from_os_str);
         let mut config = Vec::new();
         let mut color = None;