wasm.rs

  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        let fun = self
139            .instance
140            .get_typed_func::<(i32, i32), i32, _>(&mut self.store, handle.inner())?;
141
142        // call the function, passing in the buffer and its length
143        // this should return a pointer to a (ptr, lentgh) pair
144        let arg_buffer = (arg_buffer_ptr, arg_buffer_len as i32);
145        let result_buffer = fun.call(&mut self.store, arg_buffer)?;
146        dbg!(result_buffer);
147
148        // panic!();
149        // dbg!()
150
151        // create a buffer to read the (ptr, length) pair into
152        // this is a total of 4 + 4 = 8 bytes.
153        let buffer = &mut [0; 8];
154        plugin_memory.read(&mut self.store, result_buffer as usize, buffer)?;
155
156        // use these bytes (wasm stores things little-endian)
157        // to get a pointer to the buffer and its length
158        let b = buffer;
159        let result_buffer_ptr = u32::from_le_bytes([b[0], b[1], b[2], b[3]]) as usize;
160        let result_buffer_len = u32::from_le_bytes([b[4], b[5], b[6], b[7]]) as usize;
161        let result_buffer_end = result_buffer_ptr + result_buffer_len;
162
163        dbg!(result_buffer_ptr);
164        dbg!(result_buffer_len);
165
166        // read the buffer at this point into a byte array
167        // deserialize the byte array into the provided serde type
168        let result = &plugin_memory.data(&mut self.store)[result_buffer_ptr..result_buffer_end];
169        let result = bincode::deserialize(result)?;
170
171        // // deallocate the argument buffer
172        // self.free_buffer.call(&mut self.store, arg_buffer);
173
174        return Ok(result);
175    }
176
177    fn register_handle<T: AsRef<str>>(&mut self, name: T) -> bool {
178        self.instance
179            .get_export(&mut self.store, name.as_ref())
180            .is_some()
181    }
182}