1 use crate::Wizer;
2 use crate::component::ComponentInstanceState;
3 use wasmtime::component::{
4     Component, ComponentExportIndex, Instance, Lift, WasmList, types::ComponentItem,
5 };
6 use wasmtime::{Result, Store, error::Context as _, format_err};
7 
8 #[cfg(feature = "wasmprinter")]
9 use wasmtime::ToWasmtimeResult as _;
10 
11 impl Wizer {
12     /// Same as [`Wizer::run`], except for components.
run_component<T: Send>( &self, store: &mut Store<T>, wasm: &[u8], instantiate: impl AsyncFnOnce(&mut Store<T>, &Component) -> Result<Instance>, ) -> wasmtime::Result<Vec<u8>>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 
validate_component_init_func( &self, component: &Component, ) -> wasmtime::Result<ComponentExportIndex>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 
initialize_component<T: Send>( &self, store: &mut Store<T>, instance: &Instance, index: ComponentExportIndex, ) -> wasmtime::Result<()>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 /// Implementation of [`ComponentInstanceState`] backed by Wasmtime.
79 pub 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 
86 impl<T: Send> WasmtimeWizerComponent<'_, T> {
call_func<R, R2>( &mut self, instance: &str, func: &str, use_ret: impl FnOnce(&mut Store<T>, R) -> R2, ) -> R2 where R: Lift + 'static,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 
114 impl<T: Send> ComponentInstanceState for WasmtimeWizerComponent<'_, T> {
call_func_ret_list_u8( &mut self, instance: &str, func: &str, contents: impl FnOnce(&[u8]) + Send, )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 
call_func_ret_s32(&mut self, instance: &str, func: &str) -> i32127     async fn call_func_ret_s32(&mut self, instance: &str, func: &str) -> i32 {
128         self.call_func(instance, func, |_, r| r).await
129     }
130 
call_func_ret_s64(&mut self, instance: &str, func: &str) -> i64131     async fn call_func_ret_s64(&mut self, instance: &str, func: &str) -> i64 {
132         self.call_func(instance, func, |_, r| r).await
133     }
134 
call_func_ret_f32(&mut self, instance: &str, func: &str) -> u32135     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 
call_func_ret_f64(&mut self, instance: &str, func: &str) -> u64140     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 }
145