ayaka_plugin_wasmi/
lib.rs

1//! Wasmi-based plugin backend.
2
3#![warn(missing_docs)]
4
5use ayaka_plugin::*;
6use std::{
7    collections::HashMap,
8    sync::{Arc, Mutex},
9};
10use wasmi::{core::Trap, *};
11
12unsafe fn mem_slice<'a, T: 'a>(
13    store: impl Into<StoreContext<'a, T>>,
14    memory: &Memory,
15    start: i32,
16    len: i32,
17) -> &'a [u8] {
18    memory
19        .data(store)
20        .get_unchecked(start as usize..)
21        .get_unchecked(..len as usize)
22}
23
24unsafe fn mem_slice_mut<'a, T: 'a>(
25    store: impl Into<StoreContextMut<'a, T>>,
26    memory: &Memory,
27    start: i32,
28    len: i32,
29) -> &'a mut [u8] {
30    memory
31        .data_mut(store)
32        .get_unchecked_mut(start as usize..)
33        .get_unchecked_mut(..len as usize)
34}
35
36type HostStore = Arc<Mutex<Store<()>>>;
37
38/// A Wasmi [`Instance`].
39pub struct WasmiModule {
40    store: HostStore,
41    instance: Instance,
42    memory: Memory,
43    abi_free: TypedFunc<(i32, i32), ()>,
44    abi_alloc: TypedFunc<i32, i32>,
45}
46
47impl WasmiModule {
48    fn new(store: HostStore, module: &Module, linker: &wasmi::Linker<()>) -> Result<Self> {
49        let mut inner_store = store.lock().unwrap();
50        let instance = linker
51            .instantiate(inner_store.as_context_mut(), module)?
52            .start(inner_store.as_context_mut())?;
53        let memory = instance
54            .get_export(inner_store.as_context(), MEMORY_NAME)
55            .ok_or_else(|| anyhow!("cannot get memory"))?
56            .into_memory()
57            .ok_or_else(|| anyhow!("memory is not Memory"))?;
58        let abi_free = instance
59            .get_export(inner_store.as_context(), ABI_FREE_NAME)
60            .ok_or_else(|| anyhow!("cannot get abi_free"))?
61            .into_func()
62            .ok_or_else(|| anyhow!("abi_free is not Func"))?
63            .typed(inner_store.as_context())?;
64        let abi_alloc = instance
65            .get_export(inner_store.as_context_mut(), ABI_ALLOC_NAME)
66            .ok_or_else(|| anyhow!("cannot get abi_alloc"))?
67            .into_func()
68            .ok_or_else(|| anyhow!("abi_alloc is not Func"))?
69            .typed(inner_store.as_context())?;
70        drop(inner_store);
71        Ok(Self {
72            store,
73            instance,
74            memory,
75            abi_free,
76            abi_alloc,
77        })
78    }
79
80    fn call_impl<T>(
81        &self,
82        mut store: StoreContextMut<()>,
83        name: &str,
84        data: &[u8],
85        f: impl FnOnce(&[u8]) -> Result<T>,
86    ) -> Result<T> {
87        let func = self
88            .instance
89            .get_export(&store, name)
90            .ok_or_else(|| anyhow!("cannot get export {}", name))?
91            .into_func()
92            .ok_or_else(|| anyhow!("{} is not Func", name))?
93            .typed::<(i32, i32), u64>(&store)?;
94
95        let ptr = self.abi_alloc.call(&mut store, data.len() as i32)?;
96        unsafe {
97            mem_slice_mut(&mut store, &self.memory, ptr, data.len() as i32).copy_from_slice(data)
98        };
99
100        let res = func.call(&mut store, (data.len() as i32, ptr));
101
102        self.abi_free.call(&mut store, (ptr, data.len() as i32))?;
103
104        let res = res?;
105        let (len, res) = ((res >> 32) as i32, (res & 0xFFFFFFFF) as i32);
106
107        let res_data = unsafe { mem_slice(&store, &self.memory, res, len) };
108
109        let res_data = f(res_data);
110
111        self.abi_free.call(&mut store, (res, len))?;
112
113        let res_data = res_data?;
114        Ok(res_data)
115    }
116}
117
118impl RawModule for WasmiModule {
119    type Linker = WasmiLinker;
120
121    type LinkerHandle<'a> = WasmiLinkerHandle<'a>;
122
123    type Func = Func;
124
125    fn call<T>(&self, name: &str, data: &[u8], f: impl FnOnce(&[u8]) -> Result<T>) -> Result<T> {
126        self.call_impl(self.store.lock().unwrap().as_context_mut(), name, data, f)
127    }
128}
129
130/// A Wasmi [`Store`] with [`Linker`].
131pub struct WasmiLinker {
132    engine: Engine,
133    store: HostStore,
134    linker: wasmi::Linker<()>,
135}
136
137impl ayaka_plugin::Linker<WasmiModule> for WasmiLinker {
138    type Config = ();
139
140    fn new(_: ()) -> Result<Self> {
141        let engine = Engine::default();
142        let store = Store::new(&engine, ());
143        let linker = wasmi::Linker::new(&engine);
144        Ok(Self {
145            engine,
146            store: Arc::new(Mutex::new(store)),
147            linker,
148        })
149    }
150
151    fn create(&self, binary: &[u8]) -> Result<WasmiModule> {
152        let module = Module::new(&self.engine, binary)?;
153        let host = WasmiModule::new(self.store.clone(), &module, &self.linker)?;
154        Ok(host)
155    }
156
157    fn import(&mut self, ns: impl Into<String>, funcs: HashMap<String, Func>) -> Result<()> {
158        let ns = ns.into();
159        for (name, func) in funcs {
160            self.linker.define(&ns, &name, func)?;
161        }
162        Ok(())
163    }
164
165    fn wrap_raw(
166        &self,
167        f: impl (Fn(WasmiLinkerHandle, i32, i32) -> Result<Vec<u8>>) + Send + Sync + 'static,
168    ) -> Func {
169        Func::wrap(
170            self.store.lock().unwrap().as_context_mut(),
171            move |mut store: Caller<()>, len: i32, data: i32| unsafe {
172                let memory = store
173                    .get_export(MEMORY_NAME)
174                    .ok_or_else(|| Trap::new("cannot get memory"))?
175                    .into_memory()
176                    .ok_or_else(|| Trap::new("memory is not Memory"))?;
177                let data = {
178                    let store = store.as_context_mut();
179                    let handle = WasmiLinkerHandle { store, memory };
180                    f(handle, data, len).map_err(|e| Trap::new(e.to_string()))?
181                };
182                let abi_alloc = store
183                    .get_export(ABI_ALLOC_NAME)
184                    .ok_or_else(|| Trap::new("cannot get abi_alloc"))?
185                    .into_func()
186                    .ok_or_else(|| Trap::new("abi_alloc is not Func"))?
187                    .typed::<i32, i32>(store.as_context())
188                    .map_err(|e| Trap::new(e.to_string()))?;
189                let ptr = abi_alloc.call(store.as_context_mut(), data.len() as i32)?;
190                mem_slice_mut(store.as_context_mut(), &memory, ptr, data.len() as i32)
191                    .copy_from_slice(&data);
192                Ok(((data.len() as u64) << 32) | (ptr as u64))
193            },
194        )
195    }
196}
197
198/// A Wasmi [`StoreContextMut`].
199pub struct WasmiLinkerHandle<'a> {
200    store: StoreContextMut<'a, ()>,
201    memory: Memory,
202}
203
204impl<'a> LinkerHandle<'a, WasmiModule> for WasmiLinkerHandle<'a> {
205    fn call<T>(
206        &mut self,
207        m: &WasmiModule,
208        name: &str,
209        data: &[u8],
210        f: impl FnOnce(&[u8]) -> Result<T>,
211    ) -> Result<T> {
212        m.call_impl(self.store.as_context_mut(), name, data, f)
213    }
214
215    fn slice<T>(&self, start: i32, len: i32, f: impl FnOnce(&[u8]) -> T) -> T {
216        f(unsafe { mem_slice(self.store.as_context(), &self.memory, start, len) })
217    }
218
219    fn slice_mut<T>(&mut self, start: i32, len: i32, f: impl FnOnce(&mut [u8]) -> T) -> T {
220        f(unsafe { mem_slice_mut(self.store.as_context_mut(), &self.memory, start, len) })
221    }
222}
ayaka_plugin_wasmi/lib.rs

ayaka_plugin_wasmi/
lib.rs
