src/tools/miri/cargo-miri/src/arg.rs

   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 }