wasmtime_cli_flags/
opt.rs

1//! Support for parsing Wasmtime's `-O`, `-W`, etc "option groups"
2//!
3//! This builds up a clap-derive-like system where there's ideally a single
4//! macro `wasmtime_option_group!` which is invoked per-option which enables
5//! specifying options in a struct-like syntax where all other boilerplate about
6//! option parsing is contained exclusively within this module.
7
8use crate::{KeyValuePair, WasiNnGraph};
9use anyhow::{bail, Result};
10use clap::builder::{StringValueParser, TypedValueParser, ValueParserFactory};
11use clap::error::{Error, ErrorKind};
12use serde::de::{self, Visitor};
13use std::time::Duration;
14use std::{fmt, marker};
15
16/// Characters which can be safely ignored while parsing numeric options to wasmtime
17const IGNORED_NUMBER_CHARS: [char; 1] = ['_'];
18
19#[macro_export]
20macro_rules! wasmtime_option_group {
21    (
22        $(#[$attr:meta])*
23        pub struct $opts:ident {
24            $(
25                $(#[doc = $doc:tt])*
26                $(#[serde($serde_attr:meta)])*
27                pub $opt:ident: $container:ident<$payload:ty>,
28            )+
29
30            $(
31                #[prefixed = $prefix:tt]
32                $(#[serde($serde_attr2:meta)])*
33                $(#[doc = $prefixed_doc:tt])*
34                pub $prefixed:ident: Vec<(String, Option<String>)>,
35            )?
36        }
37        enum $option:ident {
38            ...
39        }
40    ) => {
41        #[derive(Default, Debug)]
42        $(#[$attr])*
43        pub struct $opts {
44            $(
45                $(#[serde($serde_attr)])*
46                pub $opt: $container<$payload>,
47            )+
48            $(
49                $(#[serde($serde_attr2)])*
50                pub $prefixed: Vec<(String, Option<String>)>,
51            )?
52        }
53
54        #[derive(Clone, PartialEq)]
55        #[expect(non_camel_case_types, reason = "macro-generated code")]
56        enum $option {
57            $(
58                $opt($payload),
59            )+
60            $(
61                $prefixed(String, Option<String>),
62            )?
63        }
64
65        impl $crate::opt::WasmtimeOption for $option {
66            const OPTIONS: &'static [$crate::opt::OptionDesc<$option>] = &[
67                $(
68                    $crate::opt::OptionDesc {
69                        name: $crate::opt::OptName::Name(stringify!($opt)),
70                        parse: |_, s| {
71                            Ok($option::$opt(
72                                $crate::opt::WasmtimeOptionValue::parse(s)?
73                            ))
74                        },
75                        val_help: <$payload as $crate::opt::WasmtimeOptionValue>::VAL_HELP,
76                        docs: concat!($($doc, "\n",)*),
77                    },
78                 )+
79                $(
80                    $crate::opt::OptionDesc {
81                        name: $crate::opt::OptName::Prefix($prefix),
82                        parse: |name, val| {
83                            Ok($option::$prefixed(
84                                name.to_string(),
85                                val.map(|v| v.to_string()),
86                            ))
87                        },
88                        val_help: "[=val]",
89                        docs: concat!($($prefixed_doc, "\n",)*),
90                    },
91                 )?
92            ];
93        }
94
95        impl core::fmt::Display for $option {
96            fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
97                match self {
98                    $(
99                        $option::$opt(val) => {
100                            write!(f, "{}=", stringify!($opt).replace('_', "-"))?;
101                            $crate::opt::WasmtimeOptionValue::display(val, f)
102                        }
103                    )+
104                    $(
105                        $option::$prefixed(key, val) => {
106                            write!(f, "{}-{key}", stringify!($prefixed))?;
107                            if let Some(val) = val {
108                                write!(f, "={val}")?;
109                            }
110                            Ok(())
111                        }
112                    )?
113                }
114            }
115        }
116
117        impl $opts {
118            fn configure_with(&mut self, opts: &[$crate::opt::CommaSeparated<$option>]) {
119                for opt in opts.iter().flat_map(|o| o.0.iter()) {
120                    match opt {
121                        $(
122                            $option::$opt(val) => {
123                                $crate::opt::OptionContainer::push(&mut self.$opt, val.clone());
124                            }
125                        )+
126                        $(
127                            $option::$prefixed(key, val) => self.$prefixed.push((key.clone(), val.clone())),
128                        )?
129                    }
130                }
131            }
132
133            fn to_options(&self) -> Vec<$option> {
134                let mut ret = Vec::new();
135                $(
136                    for item in $crate::opt::OptionContainer::get(&self.$opt) {
137                        ret.push($option::$opt(item.clone()));
138                    }
139                )+
140                $(
141                    for (key,val) in self.$prefixed.iter() {
142                        ret.push($option::$prefixed(key.clone(), val.clone()));
143                    }
144                )?
145                ret
146            }
147        }
148    };
149}
150
151/// Parser registered with clap which handles parsing the `...` in `-O ...`.
152#[derive(Clone, Debug, PartialEq)]
153pub struct CommaSeparated<T>(pub Vec<T>);
154
155impl<T> ValueParserFactory for CommaSeparated<T>
156where
157    T: WasmtimeOption,
158{
159    type Parser = CommaSeparatedParser<T>;
160
161    fn value_parser() -> CommaSeparatedParser<T> {
162        CommaSeparatedParser(marker::PhantomData)
163    }
164}
165
166#[derive(Clone)]
167pub struct CommaSeparatedParser<T>(marker::PhantomData<T>);
168
169impl<T> TypedValueParser for CommaSeparatedParser<T>
170where
171    T: WasmtimeOption,
172{
173    type Value = CommaSeparated<T>;
174
175    fn parse_ref(
176        &self,
177        cmd: &clap::Command,
178        arg: Option<&clap::Arg>,
179        value: &std::ffi::OsStr,
180    ) -> Result<Self::Value, Error> {
181        let val = StringValueParser::new().parse_ref(cmd, arg, value)?;
182
183        let options = T::OPTIONS;
184        let arg = arg.expect("should always have an argument");
185        let arg_long = arg.get_long().expect("should have a long name specified");
186        let arg_short = arg.get_short().expect("should have a short name specified");
187
188        // Handle `-O help` which dumps all the `-O` options, their messages,
189        // and then exits.
190        if val == "help" {
191            let mut max = 0;
192            for d in options {
193                max = max.max(d.name.display_string().len() + d.val_help.len());
194            }
195            println!("Available {arg_long} options:\n");
196            for d in options {
197                print!(
198                    "  -{arg_short} {:>1$}",
199                    d.name.display_string(),
200                    max - d.val_help.len()
201                );
202                print!("{}", d.val_help);
203                print!(" --");
204                if val == "help" {
205                    for line in d.docs.lines().map(|s| s.trim()) {
206                        if line.is_empty() {
207                            break;
208                        }
209                        print!(" {line}");
210                    }
211                    println!();
212                } else {
213                    println!();
214                    for line in d.docs.lines().map(|s| s.trim()) {
215                        let line = line.trim();
216                        println!("        {line}");
217                    }
218                }
219            }
220            println!("\npass `-{arg_short} help-long` to see longer-form explanations");
221            std::process::exit(0);
222        }
223        if val == "help-long" {
224            println!("Available {arg_long} options:\n");
225            for d in options {
226                println!(
227                    "  -{arg_short} {}{} --",
228                    d.name.display_string(),
229                    d.val_help
230                );
231                println!();
232                for line in d.docs.lines().map(|s| s.trim()) {
233                    let line = line.trim();
234                    println!("        {line}");
235                }
236            }
237            std::process::exit(0);
238        }
239
240        let mut result = Vec::new();
241        for val in val.split(',') {
242            // Split `k=v` into `k` and `v` where `v` is optional
243            let mut iter = val.splitn(2, '=');
244            let key = iter.next().unwrap();
245            let key_val = iter.next();
246
247            // Find `key` within `T::OPTIONS`
248            let option = options
249                .iter()
250                .filter_map(|d| match d.name {
251                    OptName::Name(s) => {
252                        let s = s.replace('_', "-");
253                        if s == key {
254                            Some((d, s))
255                        } else {
256                            None
257                        }
258                    }
259                    OptName::Prefix(s) => {
260                        let name = key.strip_prefix(s)?.strip_prefix("-")?;
261                        Some((d, name.to_string()))
262                    }
263                })
264                .next();
265
266            let (desc, key) = match option {
267                Some(pair) => pair,
268                None => {
269                    let err = Error::raw(
270                        ErrorKind::InvalidValue,
271                        format!("unknown -{arg_short} / --{arg_long} option: {key}\n"),
272                    );
273                    return Err(err.with_cmd(cmd));
274                }
275            };
276
277            result.push((desc.parse)(&key, key_val).map_err(|e| {
278                Error::raw(
279                    ErrorKind::InvalidValue,
280                    format!("failed to parse -{arg_short} option `{val}`: {e:?}\n"),
281                )
282                .with_cmd(cmd)
283            })?)
284        }
285
286        Ok(CommaSeparated(result))
287    }
288}
289
290/// Helper trait used by `CommaSeparated` which contains a list of all options
291/// supported by the option group.
292pub trait WasmtimeOption: Sized + Send + Sync + Clone + 'static {
293    const OPTIONS: &'static [OptionDesc<Self>];
294}
295
296pub struct OptionDesc<T> {
297    pub name: OptName,
298    pub docs: &'static str,
299    pub parse: fn(&str, Option<&str>) -> Result<T>,
300    pub val_help: &'static str,
301}
302
303pub enum OptName {
304    /// A named option. Note that the `str` here uses `_` instead of `-` because
305    /// it's derived from Rust syntax.
306    Name(&'static str),
307
308    /// A prefixed option which strips the specified `name`, then `-`.
309    Prefix(&'static str),
310}
311
312impl OptName {
313    fn display_string(&self) -> String {
314        match self {
315            OptName::Name(s) => s.replace('_', "-"),
316            OptName::Prefix(s) => format!("{s}-<KEY>"),
317        }
318    }
319}
320
321/// A helper trait for all types of options that can be parsed. This is what
322/// actually parses the `=val` in `key=val`
323pub trait WasmtimeOptionValue: Sized {
324    /// Help text for the value to be specified.
325    const VAL_HELP: &'static str;
326
327    /// Parses the provided value, if given, returning an error on failure.
328    fn parse(val: Option<&str>) -> Result<Self>;
329
330    /// Write the value to `f` that would parse to `self`.
331    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result;
332}
333
334impl WasmtimeOptionValue for String {
335    const VAL_HELP: &'static str = "=val";
336    fn parse(val: Option<&str>) -> Result<Self> {
337        match val {
338            Some(val) => Ok(val.to_string()),
339            None => bail!("value must be specified with `key=val` syntax"),
340        }
341    }
342
343    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
344        f.write_str(self)
345    }
346}
347
348impl WasmtimeOptionValue for u32 {
349    const VAL_HELP: &'static str = "=N";
350    fn parse(val: Option<&str>) -> Result<Self> {
351        let val = String::parse(val)?.replace(IGNORED_NUMBER_CHARS, "");
352        match val.strip_prefix("0x") {
353            Some(hex) => Ok(u32::from_str_radix(hex, 16)?),
354            None => Ok(val.parse()?),
355        }
356    }
357
358    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
359        write!(f, "{self}")
360    }
361}
362
363impl WasmtimeOptionValue for u64 {
364    const VAL_HELP: &'static str = "=N";
365    fn parse(val: Option<&str>) -> Result<Self> {
366        let val = String::parse(val)?.replace(IGNORED_NUMBER_CHARS, "");
367        match val.strip_prefix("0x") {
368            Some(hex) => Ok(u64::from_str_radix(hex, 16)?),
369            None => Ok(val.parse()?),
370        }
371    }
372
373    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
374        write!(f, "{self}")
375    }
376}
377
378impl WasmtimeOptionValue for usize {
379    const VAL_HELP: &'static str = "=N";
380    fn parse(val: Option<&str>) -> Result<Self> {
381        let val = String::parse(val)?.replace(IGNORED_NUMBER_CHARS, "");
382        match val.strip_prefix("0x") {
383            Some(hex) => Ok(usize::from_str_radix(hex, 16)?),
384            None => Ok(val.parse()?),
385        }
386    }
387
388    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
389        write!(f, "{self}")
390    }
391}
392
393impl WasmtimeOptionValue for bool {
394    const VAL_HELP: &'static str = "[=y|n]";
395    fn parse(val: Option<&str>) -> Result<Self> {
396        match val {
397            None | Some("y") | Some("yes") | Some("true") => Ok(true),
398            Some("n") | Some("no") | Some("false") => Ok(false),
399            Some(s) => bail!("unknown boolean flag `{s}`, only yes,no,<nothing> accepted"),
400        }
401    }
402
403    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
404        if *self {
405            f.write_str("y")
406        } else {
407            f.write_str("n")
408        }
409    }
410}
411
412impl WasmtimeOptionValue for Duration {
413    const VAL_HELP: &'static str = "=N|Ns|Nms|..";
414    fn parse(val: Option<&str>) -> Result<Duration> {
415        let s = String::parse(val)?;
416        // assume an integer without a unit specified is a number of seconds ...
417        if let Ok(val) = s.parse() {
418            return Ok(Duration::from_secs(val));
419        }
420
421        if let Some(num) = s.strip_suffix("s") {
422            if let Ok(val) = num.parse() {
423                return Ok(Duration::from_secs(val));
424            }
425        }
426        if let Some(num) = s.strip_suffix("ms") {
427            if let Ok(val) = num.parse() {
428                return Ok(Duration::from_millis(val));
429            }
430        }
431        if let Some(num) = s.strip_suffix("us").or(s.strip_suffix("μs")) {
432            if let Ok(val) = num.parse() {
433                return Ok(Duration::from_micros(val));
434            }
435        }
436        if let Some(num) = s.strip_suffix("ns") {
437            if let Ok(val) = num.parse() {
438                return Ok(Duration::from_nanos(val));
439            }
440        }
441
442        bail!("failed to parse duration: {s}")
443    }
444
445    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
446        let subsec = self.subsec_nanos();
447        if subsec == 0 {
448            write!(f, "{}s", self.as_secs())
449        } else if subsec % 1_000 == 0 {
450            write!(f, "{}μs", self.as_micros())
451        } else if subsec % 1_000_000 == 0 {
452            write!(f, "{}ms", self.as_millis())
453        } else {
454            write!(f, "{}ns", self.as_nanos())
455        }
456    }
457}
458
459impl WasmtimeOptionValue for wasmtime::OptLevel {
460    const VAL_HELP: &'static str = "=0|1|2|s";
461    fn parse(val: Option<&str>) -> Result<Self> {
462        match String::parse(val)?.as_str() {
463            "0" => Ok(wasmtime::OptLevel::None),
464            "1" => Ok(wasmtime::OptLevel::Speed),
465            "2" => Ok(wasmtime::OptLevel::Speed),
466            "s" => Ok(wasmtime::OptLevel::SpeedAndSize),
467            other => bail!(
468                "unknown optimization level `{}`, only 0,1,2,s accepted",
469                other
470            ),
471        }
472    }
473
474    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
475        match *self {
476            wasmtime::OptLevel::None => f.write_str("0"),
477            wasmtime::OptLevel::Speed => f.write_str("2"),
478            wasmtime::OptLevel::SpeedAndSize => f.write_str("s"),
479            _ => unreachable!(),
480        }
481    }
482}
483
484impl WasmtimeOptionValue for wasmtime::RegallocAlgorithm {
485    const VAL_HELP: &'static str = "=backtracking|single-pass";
486    fn parse(val: Option<&str>) -> Result<Self> {
487        match String::parse(val)?.as_str() {
488            "backtracking" => Ok(wasmtime::RegallocAlgorithm::Backtracking),
489            "single-pass" => Ok(wasmtime::RegallocAlgorithm::SinglePass),
490            other => bail!(
491                "unknown regalloc algorithm`{}`, only backtracking,single-pass accepted",
492                other
493            ),
494        }
495    }
496
497    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
498        match *self {
499            wasmtime::RegallocAlgorithm::Backtracking => f.write_str("backtracking"),
500            wasmtime::RegallocAlgorithm::SinglePass => f.write_str("single-pass"),
501            _ => unreachable!(),
502        }
503    }
504}
505
506impl WasmtimeOptionValue for wasmtime::Strategy {
507    const VAL_HELP: &'static str = "=winch|cranelift";
508    fn parse(val: Option<&str>) -> Result<Self> {
509        match String::parse(val)?.as_str() {
510            "cranelift" => Ok(wasmtime::Strategy::Cranelift),
511            "winch" => Ok(wasmtime::Strategy::Winch),
512            other => bail!("unknown compiler `{other}` only `cranelift` and `winch` accepted",),
513        }
514    }
515
516    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
517        match *self {
518            wasmtime::Strategy::Cranelift => f.write_str("cranelift"),
519            wasmtime::Strategy::Winch => f.write_str("winch"),
520            _ => unreachable!(),
521        }
522    }
523}
524
525impl WasmtimeOptionValue for wasmtime::Collector {
526    const VAL_HELP: &'static str = "=drc|null";
527    fn parse(val: Option<&str>) -> Result<Self> {
528        match String::parse(val)?.as_str() {
529            "drc" => Ok(wasmtime::Collector::DeferredReferenceCounting),
530            "null" => Ok(wasmtime::Collector::Null),
531            other => bail!("unknown collector `{other}` only `drc` and `null` accepted",),
532        }
533    }
534
535    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
536        match *self {
537            wasmtime::Collector::DeferredReferenceCounting => f.write_str("drc"),
538            wasmtime::Collector::Null => f.write_str("null"),
539            _ => unreachable!(),
540        }
541    }
542}
543
544impl WasmtimeOptionValue for wasmtime::MpkEnabled {
545    const VAL_HELP: &'static str = "[=y|n|auto]";
546    fn parse(val: Option<&str>) -> Result<Self> {
547        match val {
548            None | Some("y") | Some("yes") | Some("true") => Ok(wasmtime::MpkEnabled::Enable),
549            Some("n") | Some("no") | Some("false") => Ok(wasmtime::MpkEnabled::Disable),
550            Some("auto") => Ok(wasmtime::MpkEnabled::Auto),
551            Some(s) => bail!("unknown mpk flag `{s}`, only yes,no,auto,<nothing> accepted"),
552        }
553    }
554
555    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
556        match *self {
557            wasmtime::MpkEnabled::Enable => f.write_str("y"),
558            wasmtime::MpkEnabled::Disable => f.write_str("n"),
559            wasmtime::MpkEnabled::Auto => f.write_str("auto"),
560        }
561    }
562}
563
564impl WasmtimeOptionValue for WasiNnGraph {
565    const VAL_HELP: &'static str = "=<format>::<dir>";
566    fn parse(val: Option<&str>) -> Result<Self> {
567        let val = String::parse(val)?;
568        let mut parts = val.splitn(2, "::");
569        Ok(WasiNnGraph {
570            format: parts.next().unwrap().to_string(),
571            dir: match parts.next() {
572                Some(part) => part.into(),
573                None => bail!("graph does not contain `::` separator for directory"),
574            },
575        })
576    }
577
578    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
579        write!(f, "{}::{}", self.format, self.dir)
580    }
581}
582
583impl WasmtimeOptionValue for KeyValuePair {
584    const VAL_HELP: &'static str = "=<name>=<val>";
585    fn parse(val: Option<&str>) -> Result<Self> {
586        let val = String::parse(val)?;
587        let mut parts = val.splitn(2, "=");
588        Ok(KeyValuePair {
589            key: parts.next().unwrap().to_string(),
590            value: match parts.next() {
591                Some(part) => part.into(),
592                None => "".to_string(),
593            },
594        })
595    }
596
597    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
598        f.write_str(&self.key)?;
599        if !self.value.is_empty() {
600            f.write_str("=")?;
601            f.write_str(&self.value)?;
602        }
603        Ok(())
604    }
605}
606
607pub trait OptionContainer<T> {
608    fn push(&mut self, val: T);
609    fn get<'a>(&'a self) -> impl Iterator<Item = &'a T>
610    where
611        T: 'a;
612}
613
614impl<T> OptionContainer<T> for Option<T> {
615    fn push(&mut self, val: T) {
616        *self = Some(val);
617    }
618    fn get<'a>(&'a self) -> impl Iterator<Item = &'a T>
619    where
620        T: 'a,
621    {
622        self.iter()
623    }
624}
625
626impl<T> OptionContainer<T> for Vec<T> {
627    fn push(&mut self, val: T) {
628        Vec::push(self, val);
629    }
630    fn get<'a>(&'a self) -> impl Iterator<Item = &'a T>
631    where
632        T: 'a,
633    {
634        self.iter()
635    }
636}
637
638// Used to parse toml values into string so that we can reuse the `WasmtimeOptionValue::parse`
639// for parsing toml values the same way we parse command line values.
640//
641// Used for wasmtime::Strategy, wasmtime::Collector, wasmtime::OptLevel, wasmtime::RegallocAlgorithm
642struct ToStringVisitor {}
643
644impl<'de> Visitor<'de> for ToStringVisitor {
645    type Value = String;
646
647    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
648        write!(formatter, "&str, u64, or i64")
649    }
650
651    fn visit_str<E>(self, s: &str) -> Result<Self::Value, E>
652    where
653        E: de::Error,
654    {
655        Ok(s.to_owned())
656    }
657
658    fn visit_u64<E>(self, v: u64) -> Result<Self::Value, E>
659    where
660        E: de::Error,
661    {
662        Ok(v.to_string())
663    }
664
665    fn visit_i64<E>(self, v: i64) -> Result<Self::Value, E>
666    where
667        E: de::Error,
668    {
669        Ok(v.to_string())
670    }
671}
672
673// Deserializer that uses the `WasmtimeOptionValue::parse` to parse toml values
674pub(crate) fn cli_parse_wrapper<'de, D, T>(deserializer: D) -> Result<Option<T>, D::Error>
675where
676    T: WasmtimeOptionValue,
677    D: serde::Deserializer<'de>,
678{
679    let to_string_visitor = ToStringVisitor {};
680    let str = deserializer.deserialize_any(to_string_visitor)?;
681
682    T::parse(Some(&str))
683        .map(Some)
684        .map_err(serde::de::Error::custom)
685}
686
687#[cfg(test)]
688mod tests {
689    use super::WasmtimeOptionValue;
690
691    #[test]
692    fn numbers_with_underscores() {
693        assert!(<u32 as WasmtimeOptionValue>::parse(Some("123")).is_ok_and(|v| v == 123));
694        assert!(<u32 as WasmtimeOptionValue>::parse(Some("1_2_3")).is_ok_and(|v| v == 123));
695    }
696}
wasmtime_cli_flags/opt.rs

wasmtime_cli_flags/
opt.rs
