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