wasmtime_wizer/component/
wasmtime.rs

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