Skip to content
This repository was archived by the owner on May 28, 2025. It is now read-only.

Commit 08e7d94

Browse files
RalfJungRalfJung
committed
avoid some string copies...
1 parent 4655382 commit 08e7d94

File tree

4 files changed

+140
-80
lines changed

4 files changed

+140
-80
lines changed

‎cargo-miri/src/arg.rs

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
//! Utilities for dealing with argument flags
2+
3+
use std::borrow::Cow;
4+
use std::env;
5+
6+
/// Determines whether a `--flag` is present.
7+
pub fn has_arg_flag(name: &str) -> bool {
8+
num_arg_flag(name) > 0
9+
}
10+
11+
/// Determines how many times a `--flag` is present.
12+
pub fn num_arg_flag(name: &str) -> usize {
13+
env::args().take_while(|val| val != "--").filter(|val| val == name).count()
14+
}
15+
16+
/// Yields all values of command line flag `name` as `Ok(arg)`, and all other arguments except
17+
/// the flag as `Err(arg)`. (The flag `name` itself is not yielded at all, only its values are.)
18+
pub struct ArgSplitFlagValue<'a, I> {
19+
args: Option<I>,
20+
name: &'a str,
21+
}
22+
23+
impl<'a, I: Iterator> ArgSplitFlagValue<'a, I> {
24+
fn new(args: I, name: &'a str) -> Self {
25+
Self { args: Some(args), name }
26+
}
27+
}
28+
29+
impl<'s, I: Iterator<Item = Cow<'s, str>>> Iterator for ArgSplitFlagValue<'_, I> {
30+
// If the original iterator was all `Owned`, then we will only ever yield `Owned`
31+
// (so `into_owned()` is cheap).
32+
type Item = Result<Cow<'s, str>, Cow<'s, str>>;
33+
34+
fn next(&mut self) -> Option<Self::Item> {
35+
let Some(args) = self.args.as_mut() else {
36+
// We already canceled this iterator.
37+
return None;
38+
};
39+
let arg = args.next()?;
40+
if arg == "--" {
41+
// Stop searching at `--`.
42+
self.args = None;
43+
return None;
44+
}
45+
// These branches cannot be merged if we want to avoid the allocation in the `Borrowed` branch.
46+
match &arg {
47+
Cow::Borrowed(arg) =>
48+
if let Some(suffix) = arg.strip_prefix(self.name) {
49+
// Strip leading `name`.
50+
if suffix.is_empty() {
51+
// This argument is exactly `name`; the next one is the value.
52+
return args.next().map(Ok);
53+
} else if let Some(suffix) = suffix.strip_prefix('=') {
54+
// This argument is `name=value`; get the value.
55+
return Some(Ok(Cow::Borrowed(suffix)));
56+
}
57+
},
58+
Cow::Owned(arg) =>
59+
if let Some(suffix) = arg.strip_prefix(self.name) {
60+
// Strip leading `name`.
61+
if suffix.is_empty() {
62+
// This argument is exactly `name`; the next one is the value.
63+
return args.next().map(Ok);
64+
} else if let Some(suffix) = suffix.strip_prefix('=') {
65+
// This argument is `name=value`; get the value. We need to do an allocation
66+
// here as a `String` cannot be subsliced (what would the lifetime be?).
67+
return Some(Ok(Cow::Owned(suffix.to_owned())));
68+
}
69+
},
70+
}
71+
Some(Err(arg))
72+
}
73+
}
74+
75+
impl<'a, I: Iterator<Item = String> + 'a> ArgSplitFlagValue<'a, I> {
76+
pub fn from_string_iter(
77+
args: I,
78+
name: &'a str,
79+
) -> impl Iterator<Item = Result<String, String>> + 'a {
80+
ArgSplitFlagValue::new(args.map(Cow::Owned), name).map(|x| {
81+
match x {
82+
Ok(Cow::Owned(s)) => Ok(s),
83+
Err(Cow::Owned(s)) => Err(s),
84+
_ => panic!("iterator converted owned to borrowed"),
85+
}
86+
})
87+
}
88+
}
89+
90+
impl<'x: 'a, 'a, I: Iterator<Item = &'x str> + 'a> ArgSplitFlagValue<'a, I> {
91+
pub fn from_str_iter(
92+
args: I,
93+
name: &'a str,
94+
) -> impl Iterator<Item = Result<&'x str, &'x str>> + 'a {
95+
ArgSplitFlagValue::new(args.map(Cow::Borrowed), name).map(|x| {
96+
match x {
97+
Ok(Cow::Borrowed(s)) => Ok(s),
98+
Err(Cow::Borrowed(s)) => Err(s),
99+
_ => panic!("iterator converted borrowed to owned"),
100+
}
101+
})
102+
}
103+
}
104+
105+
/// Yields all values of command line flag `name`.
106+
pub struct ArgFlagValueIter;
107+
108+
impl ArgFlagValueIter {
109+
pub fn from_string_iter<'a, I: Iterator<Item = String> + 'a>(
110+
args: I,
111+
name: &'a str,
112+
) -> impl Iterator<Item = String> + 'a {
113+
ArgSplitFlagValue::from_string_iter(args, name).filter_map(Result::ok)
114+
}
115+
}
116+
117+
impl ArgFlagValueIter {
118+
pub fn from_str_iter<'x: 'a, 'a, I: Iterator<Item = &'x str> + 'a>(
119+
args: I,
120+
name: &'a str,
121+
) -> impl Iterator<Item = &'x str> + 'a {
122+
ArgSplitFlagValue::from_str_iter(args, name).filter_map(Result::ok)
123+
}
124+
}
125+
126+
/// Gets the values of a `--flag`.
127+
pub fn get_arg_flag_values(name: &str) -> impl Iterator<Item = String> + '_ {
128+
ArgFlagValueIter::from_string_iter(env::args(), name)
129+
}
130+
131+
/// Gets the value of a `--flag`.
132+
pub fn get_arg_flag_value(name: &str) -> Option<String> {
133+
get_arg_flag_values(name).next()
134+
}

‎cargo-miri/src/main.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#![feature(let_else)]
22
#![allow(clippy::useless_format, clippy::derive_partial_eq_without_eq)]
33

4+
mod arg;
45
mod phases;
56
mod setup;
67
mod util;

‎cargo-miri/src/phases.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ pub fn phase_cargo_miri(mut args: impl Iterator<Item = String>) {
119119
// Forward all arguments before `--` other than `--target-dir` and its value to Cargo.
120120
// (We want to *change* the target-dir value, so we must not forward it.)
121121
let mut target_dir = None;
122-
for arg in ArgSplitFlagValue::new(&mut args, "--target-dir") {
122+
for arg in ArgSplitFlagValue::from_string_iter(&mut args, "--target-dir") {
123123
match arg {
124124
Ok(value) => {
125125
if target_dir.is_some() {
@@ -310,7 +310,8 @@ pub fn phase_rustc(mut args: impl Iterator<Item = String>, phase: RustcPhase) {
310310
let mut cmd = miri();
311311

312312
// Ensure --emit argument for a check-only build is present.
313-
if let Some(val) = ArgFlagValueIter::new(env.args.clone().into_iter(), "--emit").next()
313+
if let Some(val) =
314+
ArgFlagValueIter::from_str_iter(env.args.iter().map(|s| s as &str), "--emit").next()
314315
{
315316
// For `no_run` tests, rustdoc passes a `--emit` flag; make sure it has the right shape.
316317
assert_eq!(val, "metadata");

‎cargo-miri/src/util.rs

Lines changed: 2 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ use std::ffi::OsString;
44
use std::fmt::Write as _;
55
use std::fs::{self, File};
66
use std::io::{self, BufWriter, Read, Write};
7-
use std::iter::TakeWhile;
87
use std::ops::Not;
98
use std::path::{Path, PathBuf};
109
use std::process::Command;
@@ -13,6 +12,8 @@ use cargo_metadata::{Metadata, MetadataCommand};
1312
use rustc_version::VersionMeta;
1413
use serde::{Deserialize, Serialize};
1514

15+
pub use crate::arg::*;
16+
1617
/// The information to run a crate with the given environment.
1718
#[derive(Clone, Serialize, Deserialize)]
1819
pub struct CrateRunEnv {
@@ -74,83 +75,6 @@ pub fn show_error(msg: String) -> ! {
7475
std::process::exit(1)
7576
}
7677

77-
/// Determines whether a `--flag` is present.
78-
pub fn has_arg_flag(name: &str) -> bool {
79-
num_arg_flag(name) > 0
80-
}
81-
82-
/// Determines how many times a `--flag` is present.
83-
pub fn num_arg_flag(name: &str) -> usize {
84-
env::args().take_while(|val| val != "--").filter(|val| val == name).count()
85-
}
86-
87-
/// Yields all values of command line flag `name` as `Ok(arg)`, and all other arguments except
88-
/// the flag as `Err(arg)`. (The flag `name` itself is not yielded at all, only its values are.)
89-
pub struct ArgSplitFlagValue<'a, I> {
90-
args: TakeWhile<I, fn(&String) -> bool>,
91-
name: &'a str,
92-
}
93-
94-
impl<'a, I: Iterator<Item = String>> ArgSplitFlagValue<'a, I> {
95-
pub fn new(args: I, name: &'a str) -> Self {
96-
Self {
97-
// Stop searching at `--`.
98-
args: args.take_while(|val| val != "--"),
99-
name,
100-
}
101-
}
102-
}
103-
104-
impl<I: Iterator<Item = String>> Iterator for ArgSplitFlagValue<'_, I> {
105-
type Item = Result<String, String>;
106-
107-
fn next(&mut self) -> Option<Self::Item> {
108-
let arg = self.args.next()?;
109-
if let Some(suffix) = arg.strip_prefix(self.name) {
110-
// Strip leading `name`.
111-
if suffix.is_empty() {
112-
// This argument is exactly `name`; the next one is the value.
113-
return self.args.next().map(Ok);
114-
} else if let Some(suffix) = suffix.strip_prefix('=') {
115-
// This argument is `name=value`; get the value.
116-
return Some(Ok(suffix.to_owned()));
117-
}
118-
}
119-
Some(Err(arg))
120-
}
121-
}
122-
123-
/// Yields all values of command line flag `name`.
124-
pub struct ArgFlagValueIter<'a, I>(ArgSplitFlagValue<'a, I>);
125-
126-
impl<'a, I: Iterator<Item = String>> ArgFlagValueIter<'a, I> {
127-
pub fn new(args: I, name: &'a str) -> Self {
128-
Self(ArgSplitFlagValue::new(args, name))
129-
}
130-
}
131-
132-
impl<I: Iterator<Item = String>> Iterator for ArgFlagValueIter<'_, I> {
133-
type Item = String;
134-
135-
fn next(&mut self) -> Option<Self::Item> {
136-
loop {
137-
if let Ok(value) = self.0.next()? {
138-
return Some(value);
139-
}
140-
}
141-
}
142-
}
143-
144-
/// Gets the values of a `--flag`.
145-
pub fn get_arg_flag_values<'a>(name: &'a str) -> impl Iterator<Item = String> + 'a {
146-
ArgFlagValueIter::new(env::args(), name)
147-
}
148-
149-
/// Gets the value of a `--flag`.
150-
pub fn get_arg_flag_value(name: &str) -> Option<String> {
151-
get_arg_flag_values(name).next()
152-
}
153-
15478
/// Escapes `s` in a way that is suitable for using it as a string literal in TOML syntax.
15579
pub fn escape_for_toml(s: &str) -> String {
15680
// We want to surround this string in quotes `"`. So we first escape all quotes,

0 commit comments

Comments
 (0)