Skip to main content

wasmtime_wizer/component/
wasmtime.rs

1use crate::Wizer;
2use crate::component::ComponentInstanceState;
3use wasmtime::component::{
4    Component, ComponentExportIndex, Instance, Lift, WasmList, types::ComponentItem,
5};
6use wasmtime::{Result, Store, error::Context as _, format_err};
7
8#[cfg(feature = "wasmprinter")]
9use wasmtime::ToWasmtimeResult as _;
10
11impl Wizer {
12    /// Same as [`Wizer::run`], except for components.
13    pub async fn run_component<T: Send>(
14        &self,
15        store: &mut Store<T>,
16        wasm: &[u8],
17        instantiate: impl AsyncFnOnce(&mut Store<T>, &Component) -> Result<Instance>,
18    ) -> wasmtime::Result<Vec<u8>> {
19        let (cx, instrumented_wasm) = self.instrument_component(wasm)?;
20
21        #[cfg(feature = "wasmprinter")]
22        log::debug!(
23            "instrumented wasm: {}",
24            wasmprinter::print_bytes(&instrumented_wasm).to_wasmtime_result()?,
25        );
26
27        let engine = store.engine();
28        let component = Component::new(engine, &instrumented_wasm)
29            .context("failed to compile the Wasm component")?;
30        let index = self.validate_component_init_func(&component)?;
31
32        let instance = instantiate(store, &component).await?;
33        self.initialize_component(store, &instance, index).await?;
34        self.snapshot_component(cx, &mut WasmtimeWizerComponent { store, instance })
35            .await
36    }
37
38    fn validate_component_init_func(
39        &self,
40        component: &Component,
41    ) -> wasmtime::Result<ComponentExportIndex> {
42        let init_func = self.get_init_func();
43        let (ty, index) = component
44            .get_export(None, init_func)
45            .ok_or_else(|| format_err!("the component does export the function `{init_func}`"))?;
46
47        let ty = match ty {
48            ComponentItem::ComponentFunc(ty) => ty,
49            _ => wasmtime::bail!("the component's `{init_func}` export is not a function",),
50        };
51
52        if ty.params().len() != 0 || ty.results().len() != 0 {
53            wasmtime::bail!(
54                "the component's `{init_func}` function export does not have type `[] -> []`",
55            );
56        }
57        Ok(index)
58    }
59
60    async fn initialize_component<T: Send>(
61        &self,
62        store: &mut Store<T>,
63        instance: &Instance,
64        index: ComponentExportIndex,
65    ) -> wasmtime::Result<()> {
66        let init_func = instance
67            .get_typed_func::<(), ()>(&mut *store, index)
68            .expect("checked by `validate_init_func`");
69        init_func
70            .call_async(&mut *store, ())
71            .await
72            .with_context(|| format!("the initialization function trapped"))?;
73
74        Ok(())
75    }
76}
77
78/// Impementation of [`ComponentInstanceState`] backed by Wasmtime.
79pub struct WasmtimeWizerComponent<'a, T: 'static> {
80    /// The Wasmtime-based store that owns the `instance` field.
81    pub store: &'a mut Store<T>,
82    /// The instance that this will load state from.
83    pub instance: Instance,
84}
85
86impl<T: Send> WasmtimeWizerComponent<'_, T> {
87    async fn call_func<R, R2>(
88        &mut self,
89        instance: &str,
90        func: &str,
91        use_ret: impl FnOnce(&mut Store<T>, R) -> R2,
92    ) -> R2
93    where
94        R: Lift + 'static,
95    {
96        log::debug!("invoking {instance}#{func}");
97        let (_, instance_export) = self
98            .instance
99            .get_export(&mut *self.store, None, instance)
100            .unwrap();
101        let (_, func_export) = self
102            .instance
103            .get_export(&mut *self.store, Some(&instance_export), func)
104            .unwrap();
105        let func = self
106            .instance
107            .get_typed_func::<(), (R,)>(&mut *self.store, func_export)
108            .unwrap();
109        let ret = func.call_async(&mut *self.store, ()).await.unwrap().0;
110        use_ret(&mut *self.store, ret)
111    }
112}
113
114impl<T: Send> ComponentInstanceState for WasmtimeWizerComponent<'_, T> {
115    async fn call_func_ret_list_u8(
116        &mut self,
117        instance: &str,
118        func: &str,
119        contents: impl FnOnce(&[u8]) + Send,
120    ) {
121        self.call_func(instance, func, |store, list: WasmList<u8>| {
122            contents(list.as_le_slice(&store));
123        })
124        .await
125    }
126
127    async fn call_func_ret_s32(&mut self, instance: &str, func: &str) -> i32 {
128        self.call_func(instance, func, |_, r| r).await
129    }
130
131    async fn call_func_ret_s64(&mut self, instance: &str, func: &str) -> i64 {
132        self.call_func(instance, func, |_, r| r).await
133    }
134
135    async fn call_func_ret_f32(&mut self, instance: &str, func: &str) -> u32 {
136        self.call_func(instance, func, |_, r: f32| r.to_bits())
137            .await
138    }
139
140    async fn call_func_ret_f64(&mut self, instance: &str, func: &str) -> u64 {
141        self.call_func(instance, func, |_, r: f64| r.to_bits())
142            .await
143    }
144}