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