Auto merge of rust-lang#3621 - RalfJung:argparse, r=RalfJung

use a little arg-parsing helper for miri-script

Author: bors

src/tools/miri/miri-script/src/args.rs

@@ -0,0 +1,136 @@
+use std::env;
+use std::iter;
+
+use anyhow::{bail, Result};
+
+pub struct Args {
+    args: iter::Peekable<env::Args>,
+    /// Set to `true` once we saw a `--`.
+    terminated: bool,
+}
+
+impl Args {
+    pub fn new() -> Self {
+        let mut args = Args { args: env::args().peekable(), terminated: false };
+        args.args.next().unwrap(); // skip program name
+        args
+    }
+
+    /// Get the next argument without any interpretation.
+    pub fn next_raw(&mut self) -> Option<String> {
+        self.args.next()
+    }
+
+    /// Consume a `-$f` flag if present.
+    pub fn get_short_flag(&mut self, flag: char) -> Result<bool> {
+        if self.terminated {
+            return Ok(false);
+        }
+        if let Some(next) = self.args.peek() {
+            if let Some(next) = next.strip_prefix("-") {
+                if let Some(next) = next.strip_prefix(flag) {
+                    if next.is_empty() {
+                        self.args.next().unwrap(); // consume this argument
+                        return Ok(true);
+                    } else {
+                        bail!("`-{flag}` followed by value");
+                    }
+                }
+            }
+        }
+        Ok(false)
+    }
+
+    /// Consume a `--$name` flag if present.
+    pub fn get_long_flag(&mut self, name: &str) -> Result<bool> {
+        if self.terminated {
+            return Ok(false);
+        }
+        if let Some(next) = self.args.peek() {
+            if let Some(next) = next.strip_prefix("--") {
+                if next == name {
+                    self.args.next().unwrap(); // consume this argument
+                    return Ok(true);
+                }
+            }
+        }
+        Ok(false)
+    }
+
+    /// Consume a `--$name val` or `--$name=val` option if present.
+    pub fn get_long_opt(&mut self, name: &str) -> Result<Option<String>> {
+        assert!(!name.is_empty());
+        if self.terminated {
+            return Ok(None);
+        }
+        let Some(next) = self.args.peek() else { return Ok(None) };
+        let Some(next) = next.strip_prefix("--") else { return Ok(None) };
+        let Some(next) = next.strip_prefix(name) else { return Ok(None) };
+        // Starts with `--flag`.
+        Ok(if let Some(val) = next.strip_prefix("=") {
+            // `--flag=val` form
+            let val = val.into();
+            self.args.next().unwrap(); // consume this argument
+            Some(val)
+        } else if next.is_empty() {
+            // `--flag val` form
+            self.args.next().unwrap(); // consume this argument
+            let Some(val) = self.args.next() else { bail!("`--{name}` not followed by value") };
+            Some(val)
+        } else {
+            // Some unrelated flag, like `--flag-more` or so.
+            None
+        })
+    }
+
+    /// Consume a `--$name=val` or `--$name` option if present; the latter
+    /// produces a default value. (`--$name val` is *not* accepted for this form
+    /// of argument, it understands `val` already as the next argument!)
+    pub fn get_long_opt_with_default(
+        &mut self,
+        name: &str,
+        default: &str,
+    ) -> Result<Option<String>> {
+        assert!(!name.is_empty());
+        if self.terminated {
+            return Ok(None);
+        }
+        let Some(next) = self.args.peek() else { return Ok(None) };
+        let Some(next) = next.strip_prefix("--") else { return Ok(None) };
+        let Some(next) = next.strip_prefix(name) else { return Ok(None) };
+        // Starts with `--flag`.
+        Ok(if let Some(val) = next.strip_prefix("=") {
+            // `--flag=val` form
+            let val = val.into();
+            self.args.next().unwrap(); // consume this argument
+            Some(val)
+        } else if next.is_empty() {
+            // `--flag` form
+            self.args.next().unwrap(); // consume this argument
+            Some(default.into())
+        } else {
+            // Some unrelated flag, like `--flag-more` or so.
+            None
+        })
+    }
+
+    /// Returns the next free argument or uninterpreted flag, or `None` if there are no more
+    /// arguments left. `--` is returned as well, but it is interpreted in the sense that no more
+    /// flags will be parsed after this.
+    pub fn get_other(&mut self) -> Option<String> {
+        if self.terminated {
+            return self.args.next();
+        }
+        let next = self.args.next()?;
+        if next == "--" {
+            self.terminated = true; // don't parse any more flags
+            // This is where our parser is special, we do yield the `--`.
+        }
+        Some(next)
+    }
+
+    /// Return the rest of the aguments entirely unparsed.
+    pub fn remainder(self) -> Vec<String> {
+        self.args.collect()
+    }
+}

src/tools/miri/miri-script/src/commands.rs

@@ -25,7 +25,11 @@ impl MiriEnv {
     /// Returns the location of the sysroot.
     ///
     /// If the target is None the sysroot will be built for the host machine.
-    fn build_miri_sysroot(&mut self, quiet: bool, target: Option<&OsStr>) -> Result<PathBuf> {
+    fn build_miri_sysroot(
+        &mut self,
+        quiet: bool,
+        target: Option<impl AsRef<OsStr>>,
+    ) -> Result<PathBuf> {
         if let Some(miri_sysroot) = self.sh.var_os("MIRI_SYSROOT") {
             // Sysroot already set, use that.
             return Ok(miri_sysroot.into());
@@ -37,14 +41,17 @@ impl MiriEnv {
         self.build(path!(self.miri_dir / "Cargo.toml"), &[], quiet)?;
         self.build(&manifest_path, &[], quiet)?;
 
-        let target_flag =
-            if let Some(target) = target { vec![OsStr::new("--target"), target] } else { vec![] };
+        let target_flag = if let Some(target) = &target {
+            vec![OsStr::new("--target"), target.as_ref()]
+        } else {
+            vec![]
+        };
         let target_flag = &target_flag;
 
         if !quiet {
             eprint!("$ cargo miri setup");
-            if let Some(target) = target {
-                eprint!(" --target {target}", target = target.to_string_lossy());
+            if let Some(target) = &target {
+                eprint!(" --target {target}", target = target.as_ref().to_string_lossy());
             }
             eprintln!();
         }
@@ -157,8 +164,8 @@ impl Command {
             Command::Build { flags } => Self::build(flags),
             Command::Check { flags } => Self::check(flags),
             Command::Test { bless, flags, target } => Self::test(bless, flags, target),
-            Command::Run { dep, verbose, many_seeds, flags } =>
-                Self::run(dep, verbose, many_seeds, flags),
+            Command::Run { dep, verbose, many_seeds, target, edition, flags } =>
+                Self::run(dep, verbose, many_seeds, target, edition, flags),
             Command::Fmt { flags } => Self::fmt(flags),
             Command::Clippy { flags } => Self::clippy(flags),
             Command::Cargo { flags } => Self::cargo(flags),
@@ -169,7 +176,7 @@ impl Command {
         }
     }
 
-    fn toolchain(flags: Vec<OsString>) -> Result<()> {
+    fn toolchain(flags: Vec<String>) -> Result<()> {
         // Make sure rustup-toolchain-install-master is installed.
         which::which("rustup-toolchain-install-master")
             .context("Please install rustup-toolchain-install-master by running 'cargo install rustup-toolchain-install-master'")?;
@@ -364,7 +371,7 @@ impl Command {
         Ok(())
     }
 
-    fn bench(target: Option<OsString>, benches: Vec<OsString>) -> Result<()> {
+    fn bench(target: Option<String>, benches: Vec<String>) -> Result<()> {
         // The hyperfine to use
         let hyperfine = env::var("HYPERFINE");
         let hyperfine = hyperfine.as_deref().unwrap_or("hyperfine -w 1 -m 5 --shell=none");
@@ -378,14 +385,14 @@ impl Command {
         let sh = Shell::new()?;
         sh.change_dir(miri_dir()?);
         let benches_dir = "bench-cargo-miri";
-        let benches = if benches.is_empty() {
+        let benches: Vec<OsString> = if benches.is_empty() {
             sh.read_dir(benches_dir)?
                 .into_iter()
                 .filter(|path| path.is_dir())
                 .map(Into::into)
                 .collect()
         } else {
-            benches.to_owned()
+            benches.into_iter().map(Into::into).collect()
         };
         let target_flag = if let Some(target) = target {
             let mut flag = OsString::from("--target=");
@@ -409,36 +416,36 @@ impl Command {
         Ok(())
     }
 
-    fn install(flags: Vec<OsString>) -> Result<()> {
+    fn install(flags: Vec<String>) -> Result<()> {
         let e = MiriEnv::new()?;
         e.install_to_sysroot(e.miri_dir.clone(), &flags)?;
         e.install_to_sysroot(path!(e.miri_dir / "cargo-miri"), &flags)?;
         Ok(())
     }
 
-    fn build(flags: Vec<OsString>) -> Result<()> {
+    fn build(flags: Vec<String>) -> Result<()> {
         let e = MiriEnv::new()?;
         e.build(path!(e.miri_dir / "Cargo.toml"), &flags, /* quiet */ false)?;
         e.build(path!(e.miri_dir / "cargo-miri" / "Cargo.toml"), &flags, /* quiet */ false)?;
         Ok(())
     }
 
-    fn check(flags: Vec<OsString>) -> Result<()> {
+    fn check(flags: Vec<String>) -> Result<()> {
         let e = MiriEnv::new()?;
         e.check(path!(e.miri_dir / "Cargo.toml"), &flags)?;
         e.check(path!(e.miri_dir / "cargo-miri" / "Cargo.toml"), &flags)?;
         Ok(())
     }
 
-    fn clippy(flags: Vec<OsString>) -> Result<()> {
+    fn clippy(flags: Vec<String>) -> Result<()> {
         let e = MiriEnv::new()?;
         e.clippy(path!(e.miri_dir / "Cargo.toml"), &flags)?;
         e.clippy(path!(e.miri_dir / "cargo-miri" / "Cargo.toml"), &flags)?;
         e.clippy(path!(e.miri_dir / "miri-script" / "Cargo.toml"), &flags)?;
         Ok(())
     }
 
-    fn cargo(flags: Vec<OsString>) -> Result<()> {
+    fn cargo(flags: Vec<String>) -> Result<()> {
         let e = MiriEnv::new()?;
         let toolchain = &e.toolchain;
         // We carefully kept the working dir intact, so this will run cargo *on the workspace in the
@@ -447,7 +454,7 @@ impl Command {
         Ok(())
     }
 
-    fn test(bless: bool, mut flags: Vec<OsString>, target: Option<OsString>) -> Result<()> {
+    fn test(bless: bool, mut flags: Vec<String>, target: Option<String>) -> Result<()> {
         let mut e = MiriEnv::new()?;
 
         // Prepare a sysroot.
@@ -475,21 +482,30 @@ impl Command {
         dep: bool,
         verbose: bool,
         many_seeds: Option<Range<u32>>,
-        mut flags: Vec<OsString>,
+        target: Option<String>,
+        edition: Option<String>,
+        flags: Vec<String>,
     ) -> Result<()> {
         let mut e = MiriEnv::new()?;
-        let target = arg_flag_value(&flags, "--target");
-
-        // Scan for "--edition", set one ourselves if that flag is not present.
-        let have_edition = arg_flag_value(&flags, "--edition").is_some();
-        if !have_edition {
-            flags.push("--edition=2021".into()); // keep in sync with `tests/ui.rs`.`
+        // More flags that we will pass before `flags`
+        // (because `flags` may contain `--`).
+        let mut early_flags = Vec::<OsString>::new();
+
+        // Add target, edition to flags.
+        if let Some(target) = &target {
+            early_flags.push("--target".into());
+            early_flags.push(target.into());
+        }
+        if verbose {
+            early_flags.push("--verbose".into());
         }
+        early_flags.push("--edition".into());
+        early_flags.push(edition.as_deref().unwrap_or("2021").into());
 
-        // Prepare a sysroot, and add it to the flags.
+        // Prepare a sysroot, add it to the flags.
         let miri_sysroot = e.build_miri_sysroot(/* quiet */ !verbose, target.as_deref())?;
-        flags.push("--sysroot".into());
-        flags.push(miri_sysroot.into());
+        early_flags.push("--sysroot".into());
+        early_flags.push(miri_sysroot.into());
 
         // Compute everything needed to run the actual command. Also add MIRIFLAGS.
         let miri_manifest = path!(e.miri_dir / "Cargo.toml");
@@ -515,7 +531,7 @@ impl Command {
             };
             cmd.set_quiet(!verbose);
             // Add Miri flags
-            let cmd = cmd.args(&miri_flags).args(seed_flag).args(&flags);
+            let cmd = cmd.args(&miri_flags).args(&seed_flag).args(&early_flags).args(&flags);
             // And run the thing.
             Ok(cmd.run()?)
         };
@@ -534,7 +550,7 @@ impl Command {
         Ok(())
     }
 
-    fn fmt(flags: Vec<OsString>) -> Result<()> {
+    fn fmt(flags: Vec<String>) -> Result<()> {
         use itertools::Itertools;
 
         let e = MiriEnv::new()?;
@@ -556,6 +572,6 @@ impl Command {
             .filter_ok(|item| item.file_type().is_file())
             .map_ok(|item| item.into_path());
 
-        e.format_files(files, &e.toolchain[..], &config_path, &flags[..])
+        e.format_files(files, &e.toolchain[..], &config_path, &flags)
     }
 }