Skip to main content

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 mut run = RunCommand {
78            run: self.run,
79            argv0: None,
80            invoke: Some(if is_component {
81                format!("{}()", self.wizer.get_init_func())
82            } else {
83                self.wizer.get_init_func().to_string()
84            }),
85            module_and_args: vec![self.input.clone().into()],
86            preloads: self.preloads.clone(),
87        };
88        let engine = run.new_engine()?;
89
90        // Instrument the input wasm with wizer.
91        let (cx, main) = if is_component {
92            #[cfg(feature = "component-model")]
93            {
94                let (cx, wasm) = self.wizer.instrument_component(&wasm)?;
95                (
96                    WizerInfo::Component(cx),
97                    RunTarget::Component(wasmtime::component::Component::new(&engine, &wasm)?),
98                )
99            }
100            #[cfg(not(feature = "component-model"))]
101            unreachable!();
102        } else {
103            let (cx, wasm) = self.wizer.instrument(&wasm)?;
104            (
105                WizerInfo::Core(cx),
106                RunTarget::Core(Module::new(&engine, &wasm)?),
107            )
108        };
109
110        // Execute a rough equivalent of
111        // `wasmtime run --invoke <..> <instrumented-wasm>`
112        let (mut store, mut linker) = run.new_store_and_linker(&engine, &main)?;
113        let instance = run
114            .instantiate_and_run(&engine, &mut linker, &main, &mut store)
115            .await?;
116
117        // Use our state to capture a snapshot with Wizer and then serialize
118        // that.
119        let final_wasm = match (cx, instance) {
120            (WizerInfo::Core(cx), CliInstance::Core(instance)) => {
121                self.wizer
122                    .snapshot(
123                        cx,
124                        &mut wasmtime_wizer::WasmtimeWizer {
125                            store: &mut store,
126                            instance,
127                        },
128                    )
129                    .await?
130            }
131
132            #[cfg(feature = "component-model")]
133            (WizerInfo::Component(cx), CliInstance::Component(instance)) => {
134                self.wizer
135                    .snapshot_component(
136                        cx,
137                        &mut wasmtime_wizer::WasmtimeWizerComponent {
138                            store: &mut store,
139                            instance,
140                        },
141                    )
142                    .await?
143            }
144
145            #[cfg(feature = "component-model")]
146            (WizerInfo::Core(_) | WizerInfo::Component(_), _) => unreachable!(),
147        };
148
149        match &self.output {
150            Some(file) => fs::write(file, &final_wasm).context("failed to write output file")?,
151            None => std::io::stdout()
152                .write_all(&final_wasm)
153                .context("failed to write output to stdout")?,
154        }
155        Ok(())
156    }
157}