avoid some string copies...

This commit is contained in:
Ralf Jung 2022-08-06 21:17:54 -04:00
parent 465538245a
commit 08e7d94562
4 changed files with 140 additions and 80 deletions

134
cargo-miri/src/arg.rs Normal file
View File

@ -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()
}

View File

@ -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;

View File

@ -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 @@ fn out_filename(prefix: &str, suffix: &str) -> PathBuf {
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");

View File

@ -4,7 +4,6 @@
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 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,