avoid some string copies...

cargo-miri/src/arg.rs

@@ -0,0 +1,134 @@
+//! Utilities for dealing with argument flags
+
+use std::borrow::Cow;
+use std::env;
+
+/// Determines whether a `--flag` is present.
+pub fn has_arg_flag(name: &str) -> bool {
+    num_arg_flag(name) > 0
+}
+
+/// Determines how many times a `--flag` is present.
+pub fn num_arg_flag(name: &str) -> usize {
+    env::args().take_while(|val| val != "--").filter(|val| val == name).count()
+}
+
+/// Yields all values of command line flag `name` as `Ok(arg)`, and all other arguments except
+/// the flag as `Err(arg)`. (The flag `name` itself is not yielded at all, only its values are.)
+pub struct ArgSplitFlagValue<'a, I> {
+    args: Option<I>,
+    name: &'a str,
+}
+
+impl<'a, I: Iterator> ArgSplitFlagValue<'a, I> {
+    fn new(args: I, name: &'a str) -> Self {
+        Self { args: Some(args), name }
+    }
+}
+
+impl<'s, I: Iterator<Item = Cow<'s, str>>> Iterator for ArgSplitFlagValue<'_, I> {
+    // If the original iterator was all `Owned`, then we will only ever yield `Owned`
+    // (so `into_owned()` is cheap).
+    type Item = Result<Cow<'s, str>, Cow<'s, str>>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let Some(args) = self.args.as_mut() else {
+            // We already canceled this iterator.
+            return None;
+        };
+        let arg = args.next()?;
+        if arg == "--" {
+            // Stop searching at `--`.
+            self.args = None;
+            return None;
+        }
+        // These branches cannot be merged if we want to avoid the allocation in the `Borrowed` branch.
+        match &arg {
+            Cow::Borrowed(arg) =>
+                if let Some(suffix) = arg.strip_prefix(self.name) {
+                    // Strip leading `name`.
+                    if suffix.is_empty() {
+                        // This argument is exactly `name`; the next one is the value.
+                        return args.next().map(Ok);
+                    } else if let Some(suffix) = suffix.strip_prefix('=') {
+                        // This argument is `name=value`; get the value.
+                        return Some(Ok(Cow::Borrowed(suffix)));
+                    }
+                },
+            Cow::Owned(arg) =>
+                if let Some(suffix) = arg.strip_prefix(self.name) {
+                    // Strip leading `name`.
+                    if suffix.is_empty() {
+                        // This argument is exactly `name`; the next one is the value.
+                        return args.next().map(Ok);
+                    } else if let Some(suffix) = suffix.strip_prefix('=') {
+                        // This argument is `name=value`; get the value. We need to do an allocation
+                        // here as a `String` cannot be subsliced (what would the lifetime be?).
+                        return Some(Ok(Cow::Owned(suffix.to_owned())));
+                    }
+                },
+        }
+        Some(Err(arg))
+    }
+}
+
+impl<'a, I: Iterator<Item = String> + 'a> ArgSplitFlagValue<'a, I> {
+    pub fn from_string_iter(
+        args: I,
+        name: &'a str,
+    ) -> impl Iterator<Item = Result<String, String>> + 'a {
+        ArgSplitFlagValue::new(args.map(Cow::Owned), name).map(|x| {
+            match x {
+                Ok(Cow::Owned(s)) => Ok(s),
+                Err(Cow::Owned(s)) => Err(s),
+                _ => panic!("iterator converted owned to borrowed"),
+            }
+        })
+    }
+}
+
+impl<'x: 'a, 'a, I: Iterator<Item = &'x str> + 'a> ArgSplitFlagValue<'a, I> {
+    pub fn from_str_iter(
+        args: I,
+        name: &'a str,
+    ) -> impl Iterator<Item = Result<&'x str, &'x str>> + 'a {
+        ArgSplitFlagValue::new(args.map(Cow::Borrowed), name).map(|x| {
+            match x {
+                Ok(Cow::Borrowed(s)) => Ok(s),
+                Err(Cow::Borrowed(s)) => Err(s),
+                _ => panic!("iterator converted borrowed to owned"),
+            }
+        })
+    }
+}
+
+/// Yields all values of command line flag `name`.
+pub struct ArgFlagValueIter;
+
+impl ArgFlagValueIter {
+    pub fn from_string_iter<'a, I: Iterator<Item = String> + 'a>(
+        args: I,
+        name: &'a str,
+    ) -> impl Iterator<Item = String> + 'a {
+        ArgSplitFlagValue::from_string_iter(args, name).filter_map(Result::ok)
+    }
+}
+
+impl ArgFlagValueIter {
+    pub fn from_str_iter<'x: 'a, 'a, I: Iterator<Item = &'x str> + 'a>(
+        args: I,
+        name: &'a str,
+    ) -> impl Iterator<Item = &'x str> + 'a {
+        ArgSplitFlagValue::from_str_iter(args, name).filter_map(Result::ok)
+    }
+}
+
+/// Gets the values of a `--flag`.
+pub fn get_arg_flag_values(name: &str) -> impl Iterator<Item = String> + '_ {
+    ArgFlagValueIter::from_string_iter(env::args(), name)
+}
+
+/// Gets the value of a `--flag`.
+pub fn get_arg_flag_value(name: &str) -> Option<String> {
+    get_arg_flag_values(name).next()
+}

cargo-miri/src/main.rs

@@ -1,6 +1,7 @@
 #![feature(let_else)]
 #![allow(clippy::useless_format, clippy::derive_partial_eq_without_eq)]
 
+mod arg;
 mod phases;
 mod setup;
 mod util;

cargo-miri/src/phases.rs

@@ -119,7 +119,7 @@ pub fn phase_cargo_miri(mut args: impl Iterator<Item = String>) {
     // Forward all arguments before `--` other than `--target-dir` and its value to Cargo.
     // (We want to *change* the target-dir value, so we must not forward it.)
     let mut target_dir = None;
-    for arg in ArgSplitFlagValue::new(&mut args, "--target-dir") {
+    for arg in ArgSplitFlagValue::from_string_iter(&mut args, "--target-dir") {
         match arg {
             Ok(value) => {
                 if target_dir.is_some() {
@@ -310,7 +310,8 @@ pub fn phase_rustc(mut args: impl Iterator<Item = String>, phase: RustcPhase) {
             let mut cmd = miri();
 
             // Ensure --emit argument for a check-only build is present.
-            if let Some(val) = ArgFlagValueIter::new(env.args.clone().into_iter(), "--emit").next()
+            if let Some(val) =
+                ArgFlagValueIter::from_str_iter(env.args.iter().map(|s| s as &str), "--emit").next()
             {
                 // For `no_run` tests, rustdoc passes a `--emit` flag; make sure it has the right shape.
                 assert_eq!(val, "metadata");

cargo-miri/src/util.rs

@@ -4,7 +4,6 @@ use std::ffi::OsString;
 use std::fmt::Write as _;
 use std::fs::{self, File};
 use std::io::{self, BufWriter, Read, Write};
-use std::iter::TakeWhile;
 use std::ops::Not;
 use std::path::{Path, PathBuf};
 use std::process::Command;
@@ -13,6 +12,8 @@ use cargo_metadata::{Metadata, MetadataCommand};
 use rustc_version::VersionMeta;
 use serde::{Deserialize, Serialize};
 
+pub use crate::arg::*;
+
 /// The information to run a crate with the given environment.
 #[derive(Clone, Serialize, Deserialize)]
 pub struct CrateRunEnv {
@@ -74,83 +75,6 @@ pub fn show_error(msg: String) -> ! {
     std::process::exit(1)
 }
 
-/// Determines whether a `--flag` is present.
-pub fn has_arg_flag(name: &str) -> bool {
-    num_arg_flag(name) > 0
-}
-
-/// Determines how many times a `--flag` is present.
-pub fn num_arg_flag(name: &str) -> usize {
-    env::args().take_while(|val| val != "--").filter(|val| val == name).count()
-}
-
-/// Yields all values of command line flag `name` as `Ok(arg)`, and all other arguments except
-/// the flag as `Err(arg)`. (The flag `name` itself is not yielded at all, only its values are.)
-pub struct ArgSplitFlagValue<'a, I> {
-    args: TakeWhile<I, fn(&String) -> bool>,
-    name: &'a str,
-}
-
-impl<'a, I: Iterator<Item = String>> ArgSplitFlagValue<'a, I> {
-    pub fn new(args: I, name: &'a str) -> Self {
-        Self {
-            // Stop searching at `--`.
-            args: args.take_while(|val| val != "--"),
-            name,
-        }
-    }
-}
-
-impl<I: Iterator<Item = String>> Iterator for ArgSplitFlagValue<'_, I> {
-    type Item = Result<String, String>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        let arg = self.args.next()?;
-        if let Some(suffix) = arg.strip_prefix(self.name) {
-            // Strip leading `name`.
-            if suffix.is_empty() {
-                // This argument is exactly `name`; the next one is the value.
-                return self.args.next().map(Ok);
-            } else if let Some(suffix) = suffix.strip_prefix('=') {
-                // This argument is `name=value`; get the value.
-                return Some(Ok(suffix.to_owned()));
-            }
-        }
-        Some(Err(arg))
-    }
-}
-
-/// Yields all values of command line flag `name`.
-pub struct ArgFlagValueIter<'a, I>(ArgSplitFlagValue<'a, I>);
-
-impl<'a, I: Iterator<Item = String>> ArgFlagValueIter<'a, I> {
-    pub fn new(args: I, name: &'a str) -> Self {
-        Self(ArgSplitFlagValue::new(args, name))
-    }
-}
-
-impl<I: Iterator<Item = String>> Iterator for ArgFlagValueIter<'_, I> {
-    type Item = String;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        loop {
-            if let Ok(value) = self.0.next()? {
-                return Some(value);
-            }
-        }
-    }
-}
-
-/// Gets the values of a `--flag`.
-pub fn get_arg_flag_values<'a>(name: &'a str) -> impl Iterator<Item = String> + 'a {
-    ArgFlagValueIter::new(env::args(), name)
-}
-
-/// Gets the value of a `--flag`.
-pub fn get_arg_flag_value(name: &str) -> Option<String> {
-    get_arg_flag_values(name).next()
-}
-
 /// Escapes `s` in a way that is suitable for using it as a string literal in TOML syntax.
 pub fn escape_for_toml(s: &str) -> String {
     // We want to surround this string in quotes `"`. So we first escape all quotes,