wasmtime_wizer/
wasmtime.rs

1use crate::{InstanceState, SnapshotVal, Wizer};
2use wasmparser::ValType;
3use wasmtime::error::Context;
4
5use wasmtime::{Extern, Instance, Module, Result, Store, Val};
6
7impl Wizer {
8    /// Initialize the given Wasm, snapshot it, and return the serialized
9    /// snapshot as a new, pre-initialized Wasm module.
10    pub async fn run<T: Send>(
11        &self,
12        store: &mut Store<T>,
13        wasm: &[u8],
14        instantiate: impl AsyncFnOnce(&mut Store<T>, &Module) -> Result<wasmtime::Instance>,
15    ) -> wasmtime::Result<Vec<u8>> {
16        let (cx, instrumented_wasm) = self.instrument(wasm)?;
17
18        let engine = store.engine();
19        let module = wasmtime::Module::new(engine, &instrumented_wasm)
20            .context("failed to compile the Wasm module")?;
21        self.validate_init_func(&module)?;
22
23        let instance = instantiate(store, &module).await?;
24        self.initialize(store, &instance).await?;
25        self.snapshot(cx, &mut WasmtimeWizer { store, instance })
26            .await
27    }
28
29    /// Check that the module exports an initialization function, and that the
30    /// function has the correct type.
31    fn validate_init_func(&self, module: &wasmtime::Module) -> wasmtime::Result<()> {
32        log::debug!("Validating the exported initialization function");
33        match module.get_export(self.get_init_func()) {
34            Some(wasmtime::ExternType::Func(func_ty)) => {
35                if func_ty.params().len() != 0 || func_ty.results().len() != 0 {
36                    wasmtime::bail!(
37                        "the Wasm module's `{}` function export does not have type `[] -> []`",
38                        self.get_init_func()
39                    );
40                }
41            }
42            Some(_) => wasmtime::bail!(
43                "the Wasm module's `{}` export is not a function",
44                self.get_init_func()
45            ),
46            None => wasmtime::bail!(
47                "the Wasm module does not have a `{}` export",
48                self.get_init_func()
49            ),
50        }
51        Ok(())
52    }
53
54    /// Instantiate the module and call its initialization function.
55    async fn initialize<T: Send>(
56        &self,
57        store: &mut Store<T>,
58        instance: &wasmtime::Instance,
59    ) -> wasmtime::Result<()> {
60        log::debug!("Calling the initialization function");
61
62        if let Some(export) = instance.get_export(&mut *store, "_initialize") {
63            if let Extern::Func(func) = export {
64                func.typed::<(), ()>(&store)?
65                    .call_async(&mut *store, ())
66                    .await
67                    .context("calling the Reactor initialization function")?;
68
69                if self.get_init_func() == "_initialize" {
70                    // Don't run `_initialize` twice if the it was explicitly
71                    // requested as the init function.
72                    return Ok(());
73                }
74            }
75        }
76
77        let init_func = instance
78            .get_typed_func::<(), ()>(&mut *store, self.get_init_func())
79            .expect("checked by `validate_init_func`");
80        init_func
81            .call_async(&mut *store, ())
82            .await
83            .with_context(|| format!("the `{}` function trapped", self.get_init_func()))?;
84
85        Ok(())
86    }
87}
88
89/// Implementation of [`InstanceState`] backed by Wasmtime.
90pub struct WasmtimeWizer<'a, T: 'static> {
91    /// The Wasmtime-based store that owns the `instance` field.
92    pub store: &'a mut Store<T>,
93    /// The instance that this will load state from.
94    pub instance: Instance,
95}
96
97impl<T: Send> InstanceState for WasmtimeWizer<'_, T> {
98    async fn global_get(&mut self, name: &str, _: ValType) -> SnapshotVal {
99        let global = self.instance.get_global(&mut *self.store, name).unwrap();
100        match global.get(&mut *self.store) {
101            Val::I32(x) => SnapshotVal::I32(x),
102            Val::I64(x) => SnapshotVal::I64(x),
103            Val::F32(x) => SnapshotVal::F32(x),
104            Val::F64(x) => SnapshotVal::F64(x),
105            Val::V128(x) => SnapshotVal::V128(x.as_u128()),
106            _ => panic!("unsupported global value type"),
107        }
108    }
109
110    async fn memory_contents(&mut self, name: &str, contents: impl FnOnce(&[u8]) + Send) {
111        let memory = self.instance.get_memory(&mut *self.store, name).unwrap();
112        contents(memory.data(&self.store))
113    }
114}
wasmtime_wizer/wasmtime.rs

wasmtime_wizer/
wasmtime.rs
