wasmtime_cli/commands/
wizer.rs

1use crate::commands::run::{CliInstance, Preloads, RunCommand};
2use crate::common::{RunCommon, RunTarget};
3use std::fs;
4use std::io::{self, Read, Write};
5use std::path::PathBuf;
6use wasmtime::{Module, Result, error::Context as _};
7use wasmtime_wizer::Wizer;
8
9#[derive(clap::Parser)]
10#[expect(missing_docs, reason = "inheriting wizer's docs")]
11pub struct WizerCommand {
12    #[command(flatten)]
13    run: RunCommon,
14
15    #[command(flatten)]
16    wizer: Wizer,
17
18    /// The input Wasm module's file path.
19    input: PathBuf,
20
21    #[command(flatten)]
22    preloads: Preloads,
23
24    /// The file path to write the output Wasm module to.
25    ///
26    /// If not specified, then `stdout` is used.
27    #[arg(short = 'o', long)]
28    output: Option<PathBuf>,
29}
30
31enum WizerInfo<'a> {
32    Core(wasmtime_wizer::ModuleContext<'a>),
33    #[cfg(feature = "component-model")]
34    Component(wasmtime_wizer::ComponentContext<'a>),
35}
36
37impl WizerCommand {
38    /// Runs the command.
39    pub fn execute(mut self) -> Result<()> {
40        self.run.common.init_logging()?;
41        let runtime = tokio::runtime::Builder::new_multi_thread()
42            .enable_time()
43            .enable_io()
44            .build()?;
45        runtime.block_on(self.execute_async())
46    }
47
48    async fn execute_async(mut self) -> Result<()> {
49        // By default use deterministic relaxed simd operations to guarantee
50        // that if relaxed simd operations are used in a module that they always
51        // produce the same result.
52        if self.run.common.wasm.relaxed_simd_deterministic.is_none() {
53            self.run.common.wasm.relaxed_simd_deterministic = Some(true);
54        }
55
56        // Don't provide any WASI imports by default to wizened components. The
57        // `run` command provides the "cli" world as a default so turn that off
58        // here if the command line flags don't otherwise say what to do.
59        if self.run.common.wasi.cli.is_none() {
60            self.run.common.wasi.cli = Some(false);
61        }
62
63        // Read the input wasm, possibly from stdin.
64        let mut wasm = Vec::new();
65        if self.input.to_str() == Some("-") {
66            io::stdin()
67                .read_to_end(&mut wasm)
68                .context("failed to read input Wasm module from stdin")?;
69        } else {
70            wasm = fs::read(&self.input).context("failed to read input Wasm module")?;
71        }
72
73        #[cfg(feature = "wat")]
74        let wasm = wat::parse_bytes(&wasm)?;
75        let is_component = wasmparser::Parser::is_component(&wasm);
76
77        let init_func = self.wizer.get_init_func();
78
79        let mut run = RunCommand {
80            run: self.run,
81            argv0: None,
82            invoke: Some(if is_component {
83                if !init_func.contains(')') {
84                    format!("{init_func}()")
85                } else {
86                    init_func.to_string()
87                }
88            } else {
89                init_func.to_string()
90            }),
91            module_and_args: vec![self.input.clone().into()],
92            preloads: self.preloads.clone(),
93            module_bytes: None,
94        };
95        let engine = run.new_engine()?;
96
97        // Instrument the input wasm with wizer.
98        let (cx, main) = if is_component {
99            #[cfg(feature = "component-model")]
100            {
101                let (cx, wasm) = self.wizer.instrument_component(&wasm)?;
102                (
103                    WizerInfo::Component(cx),
104                    RunTarget::Component(wasmtime::component::Component::new(&engine, &wasm)?),
105                )
106            }
107            #[cfg(not(feature = "component-model"))]
108            unreachable!();
109        } else {
110            let (cx, wasm) = self.wizer.instrument(&wasm)?;
111            (
112                WizerInfo::Core(cx),
113                RunTarget::Core(Module::new(&engine, &wasm)?),
114            )
115        };
116
117        // Execute a rough equivalent of
118        // `wasmtime run --invoke <..> <instrumented-wasm>`
119        let (mut store, mut linker) = run.new_store_and_linker(&engine, &main)?;
120        let instance = run
121            .instantiate_and_run(&engine, &mut linker, &main, &mut store)
122            .await?;
123
124        // Use our state to capture a snapshot with Wizer and then serialize
125        // that.
126        let final_wasm = match (cx, instance) {
127            (WizerInfo::Core(cx), CliInstance::Core(instance)) => {
128                self.wizer
129                    .snapshot(
130                        cx,
131                        &mut wasmtime_wizer::WasmtimeWizer {
132                            store: &mut store,
133                            instance,
134                        },
135                    )
136                    .await?
137            }
138
139            #[cfg(feature = "component-model")]
140            (WizerInfo::Component(cx), CliInstance::Component(instance)) => {
141                self.wizer
142                    .snapshot_component(
143                        cx,
144                        &mut wasmtime_wizer::WasmtimeWizerComponent {
145                            store: &mut store,
146                            instance,
147                        },
148                    )
149                    .await?
150            }
151
152            #[cfg(feature = "component-model")]
153            (WizerInfo::Core(_) | WizerInfo::Component(_), _) => unreachable!(),
154        };
155
156        match &self.output {
157            Some(file) => fs::write(file, &final_wasm).context("failed to write output file")?,
158            None => std::io::stdout()
159                .write_all(&final_wasm)
160                .context("failed to write output to stdout")?,
161        }
162        Ok(())
163    }
164}
wasmtime_cli/commands/wizer.rs

wasmtime_cli/commands/
wizer.rs
