1use std::collections::HashMap;
2
3use anyhow::anyhow;
4
5use wasmtime::{Engine, Func, Instance, Memory, MemoryType, Module, Store, TypedFunc};
6
7use crate::*;
8
9pub struct Wasm<T> {
10 engine: Engine,
11 module: Module,
12 store: Store<T>,
13 instance: Instance,
14 alloc_buffer: TypedFunc<i32, i32>,
15 // free_buffer: TypedFunc<(i32, i32), ()>,
16}
17
18pub struct WasmPlugin<T> {
19 pub source_bytes: Vec<u8>,
20 pub store_data: T,
21}
22
23impl<T> Wasm<T> {
24 pub fn dump_memory(data: &[u8]) {
25 for (i, byte) in data.iter().enumerate() {
26 if i % 32 == 0 {
27 println!();
28 }
29 if i % 4 == 0 {
30 print!("|");
31 }
32 if *byte == 0 {
33 print!("__")
34 } else {
35 print!("{:02x}", byte);
36 }
37 }
38 println!();
39 }
40}
41
42impl<S> Runtime for Wasm<S> {
43 type Plugin = WasmPlugin<S>;
44 type Error = anyhow::Error;
45
46 fn init(plugin: WasmPlugin<S>) -> Result<Self, Self::Error> {
47 let engine = Engine::default();
48 let module = Module::new(&engine, plugin.source_bytes)?;
49 let mut store: Store<S> = Store::new(&engine, plugin.store_data);
50 let instance = Instance::new(&mut store, &module, &[])?;
51
52 let alloc_buffer = instance.get_typed_func(&mut store, "__alloc_buffer")?;
53 // let free_buffer = instance.get_typed_func(&mut store, "__free_buffer")?;
54
55 Ok(Wasm {
56 engine,
57 module,
58 store,
59 instance,
60 alloc_buffer,
61 // free_buffer,
62 })
63 }
64
65 fn constant<T: DeserializeOwned>(&mut self, handle: &Handle) -> Result<T, Self::Error> {
66 let export = self
67 .instance
68 .get_export(&mut self.store, handle.inner())
69 .ok_or_else(|| anyhow!("Could not get export"))?;
70
71 todo!()
72 }
73
74 // So this call function is kinda a dance, I figured it'd be a good idea to document it.
75 // the high level is we take a serde type, serialize it to a byte array,
76 // (we're doing this using bincode for now)
77 // then toss that byte array into webassembly.
78 // webassembly grabs that byte array, does some magic,
79 // and serializes the result into yet another byte array.
80 // we then grab *that* result byte array and deserialize it into a result.
81 //
82 // phew...
83 //
84 // now the problem is, webassambly doesn't support buffers.
85 // only really like i32s, that's it (yeah, it's sad. Not even unsigned!)
86 // (ok, I'm exaggerating a bit).
87 //
88 // the Wasm function that this calls must have a very specific signature:
89 //
90 // fn(pointer to byte array: i32, length of byte array: i32)
91 // -> pointer to (
92 // pointer to byte_array: i32,
93 // length of byte array: i32,
94 // ): i32
95 //
96 // This pair `(pointer to byte array, length of byte array)` is called a `Buffer`
97 // and can be found in the cargo_test plugin.
98 //
99 // so on the wasm side, we grab the two parameters to the function,
100 // stuff them into a `Buffer`,
101 // and then pray to the `unsafe` Rust gods above that a valid byte array pops out.
102 //
103 // On the flip side, when returning from a wasm function,
104 // we convert whatever serialized result we get into byte array,
105 // which we stuff into a Buffer and allocate on the heap,
106 // which pointer to we then return.
107 // Note the double indirection!
108 //
109 // So when returning from a function, we actually leak memory *twice*:
110 //
111 // 1) once when we leak the byte array
112 // 2) again when we leak the allocated `Buffer`
113 //
114 // This isn't a problem because Wasm stops executing after the function returns,
115 // so the heap is still valid for our inspection when we want to pull things out.
116
117 // TODO: dont' use as for conversions
118 fn call<A: Serialize, R: DeserializeOwned>(
119 &mut self,
120 handle: &Handle,
121 arg: A,
122 ) -> Result<R, Self::Error> {
123 // serialize the argument using bincode
124 let arg = bincode::serialize(&arg)?;
125 let arg_buffer_len = arg.len();
126
127 // allocate a buffer and write the argument to that buffer
128 let arg_buffer_ptr = self
129 .alloc_buffer
130 .call(&mut self.store, arg_buffer_len as i32)?;
131 let plugin_memory = self
132 .instance
133 .get_memory(&mut self.store, "memory")
134 .ok_or_else(|| anyhow!("Could not grab slice of plugin memory"))?;
135 plugin_memory.write(&mut self.store, arg_buffer_ptr as usize, &arg)?;
136
137 // get the webassembly function we want to actually call
138 // TODO: precompute handle
139 let fun_name = format!("__{}", handle.inner());
140 let fun = self
141 .instance
142 .get_typed_func::<(i32, i32), i32, _>(&mut self.store, &fun_name)?;
143
144 // call the function, passing in the buffer and its length
145 // this should return a pointer to a (ptr, lentgh) pair
146 let arg_buffer = (arg_buffer_ptr, arg_buffer_len as i32);
147 let result_buffer = fun.call(&mut self.store, arg_buffer)?;
148
149 // create a buffer to read the (ptr, length) pair into
150 // this is a total of 4 + 4 = 8 bytes.
151 let buffer = &mut [0; 8];
152 plugin_memory.read(&mut self.store, result_buffer as usize, buffer)?;
153
154 // use these bytes (wasm stores things little-endian)
155 // to get a pointer to the buffer and its length
156 let b = buffer;
157 let result_buffer_ptr = u32::from_le_bytes([b[0], b[1], b[2], b[3]]) as usize;
158 let result_buffer_len = u32::from_le_bytes([b[4], b[5], b[6], b[7]]) as usize;
159 let result_buffer_end = result_buffer_ptr + result_buffer_len;
160
161 // read the buffer at this point into a byte array
162 // deserialize the byte array into the provided serde type
163 let result = &plugin_memory.data(&mut self.store)[result_buffer_ptr..result_buffer_end];
164 let result = bincode::deserialize(result)?;
165
166 // TODO: this is handled wasm-side, but I'd like to double-check
167 // // deallocate the argument buffer
168 // self.free_buffer.call(&mut self.store, arg_buffer);
169
170 return Ok(result);
171 }
172
173 fn register_handle<T: AsRef<str>>(&mut self, name: T) -> bool {
174 self.instance
175 .get_export(&mut self.store, name.as_ref())
176 .is_some()
177 }
178}