diff --git a/README.md b/README.md index 2ca3ca7..e653bc9 100644 --- a/README.md +++ b/README.md @@ -126,11 +126,11 @@ extern crate proc_macro; use proc_macro::TokenStream; -static WASM: &[u8] = include_bytes!("my_macro.wasm"); +static WASM: watt::Instance = watt::Instance::new(include_bytes!("my_macro.wasm")); #[proc_macro] pub fn my_macro(input: TokenStream) -> TokenStream { - watt::proc_macro("my_macro", input, WASM) + WASM.proc_macro("my_macro", input) } ``` diff --git a/demo/wa/src/lib.rs b/demo/wa/src/lib.rs index b30d773..a9b0ebf 100644 --- a/demo/wa/src/lib.rs +++ b/demo/wa/src/lib.rs @@ -2,11 +2,11 @@ extern crate proc_macro; use proc_macro::TokenStream; -static WASM: &[u8] = include_bytes! { +static WASM: watt::Instance = watt::Instance::new(include_bytes! { "../../impl/target/wasm32-unknown-unknown/release/watt_demo.wasm" -}; +}); #[proc_macro_derive(Demo)] pub fn demo(input: TokenStream) -> TokenStream { - watt::proc_macro_derive("demo", input, WASM) + WASM.proc_macro_derive("demo", input) } diff --git a/src/exec.rs b/src/exec.rs index e8179f4..e3e740d 100644 --- a/src/exec.rs +++ b/src/exec.rs @@ -1,82 +1,152 @@ use crate::data::Data; use crate::import; +use crate::Instance; use proc_macro::TokenStream; +use std::cell::RefCell; +use std::collections::hash_map::{HashMap, Entry}; -pub fn proc_macro(fun: &str, inputs: Vec, wasm: &[u8]) -> TokenStream { - _proc_macro(fun, inputs, wasm) +pub fn proc_macro(fun: &str, inputs: Vec, instance: &Instance) -> TokenStream { + _proc_macro(fun, inputs, instance) } #[cfg(jit)] -fn _proc_macro(fun: &str, inputs: Vec, wasm: &[u8]) -> TokenStream { +fn _proc_macro(fun: &str, inputs: Vec, instance: &crate::Instance) -> TokenStream { use crate::runtime::*; - let engine = Engine::new(); - let mut store = Store::new(&engine); - let module = Module::new(&store, wasm); - let imports = import::extern_vals(&module, &mut store); - let module_instance = Instance::new(&store, &module, &imports).unwrap(); - let main = module - .exports() - .iter() - .position(|p| p.name() == fun) - .unwrap(); - let exports = module_instance.exports(); - let main = exports[main].func().unwrap(); - let memory = exports.iter().filter_map(|e| e.memory()).next().unwrap(); - - let _guard = Data::guard(); - let args = Data::with(|d| { - inputs - .into_iter() - .map(|input| Val::i32(d.tokenstream.push(input) as i32)) - .collect::>() - }); - - current_memory::set(&memory, || { - let values = main.call(&args).unwrap(); - let handle = values.into_iter().next().unwrap(); - let handle = handle.as_i32().unwrap() as u32; - Data::with(|d| d.tokenstream[handle].clone()) + struct ThreadState { + _engine: Engine, + store: Store, + modules: HashMap, + instances: HashMap, + } + + std::thread_local! { + static STATE: RefCell = { + let engine = Engine::new(); + let store = Store::new(&engine); + RefCell::new(ThreadState { + _engine: engine, + store, + modules: HashMap::new(), + instances: HashMap::new(), + }) + }; + } + + impl ThreadState { + pub fn instance(&mut self, instance: &crate::Instance) -> (&Module, &Instance) { + let id = instance.id(); + let entry = match self.instances.entry(id) { + Entry::Occupied(e) => return (&self.modules[&id], &*e.into_mut()), + Entry::Vacant(v) => v, + }; + + let module = Module::new(&self.store, instance.wasm); + let imports = import::extern_vals(&module, &mut self.store); + let module_instance = Instance::new(&self.store, &module, &imports).unwrap(); + self.modules.insert(id, module); + let instance = entry.insert(module_instance); + (&self.modules[&id], &*instance) + } + } + + STATE.with(|state| { + let mut state = state.borrow_mut(); + let (module, instance) = state.instance(instance); + let main = module + .exports() + .iter() + .position(|p| p.name() == fun) + .unwrap(); + let exports = instance.exports(); + let main = exports[main].func().unwrap(); + let memory = exports.iter().filter_map(|e| e.memory()).next().unwrap(); + + let _guard = Data::guard(); + let args = Data::with(|d| { + inputs + .into_iter() + .map(|input| Val::i32(d.tokenstream.push(input) as i32)) + .collect::>() + }); + + current_memory::set(&memory, || { + let values = main.call(&args).unwrap(); + let handle = values.into_iter().next().unwrap(); + let handle = handle.as_i32().unwrap() as u32; + Data::with(|d| d.tokenstream[handle].clone()) + }) }) } #[cfg(not(jit))] -fn _proc_macro(fun: &str, inputs: Vec, wasm: &[u8]) -> TokenStream { +fn _proc_macro(fun: &str, inputs: Vec, instance: &Instance) -> TokenStream { use crate::runtime::{ - decode_module, get_export, init_store, instantiate_module, invoke_func, ExternVal, Value, + decode_module, get_export, init_store, instantiate_module, invoke_func, + runtime::ModuleInst, ExternVal, Store, Value, }; use std::io::Cursor; + use std::rc::Rc; - let cursor = Cursor::new(wasm); - let module = decode_module(cursor).unwrap(); - #[cfg(watt_debug)] - crate::debug::print_module(&module); - - let mut store = init_store(); - let extern_vals = import::extern_vals(&module, &mut store); - let module_instance = instantiate_module(&mut store, module, &extern_vals).unwrap(); - let main = match get_export(&module_instance, fun) { - Ok(ExternVal::Func(main)) => main, - _ => unimplemented!("unresolved macro: {:?}", fun), - }; + struct ThreadState { + store: Store, + instances: HashMap>, + } - let _guard = Data::guard(); - let args = Data::with(|d| { - inputs - .into_iter() - .map(|input| Value::I32(d.tokenstream.push(input))) - .collect() - }); - - let res = invoke_func(&mut store, main, args); - let values = match res { - Ok(values) => values, - Err(err) => panic!("{:?}", err), - }; - let handle = values.into_iter().next().unwrap(); - let handle = match handle { - Value::I32(handle) => handle, - _ => unimplemented!("unexpected macro return type"), - }; - Data::with(|d| d.tokenstream[handle].clone()) + std::thread_local! { + static STATE: RefCell = { + RefCell::new(ThreadState { + store: init_store(), + instances: HashMap::new(), + }) + }; + } + + impl ThreadState { + pub fn instance(&mut self, instance: &crate::Instance) -> &ModuleInst { + let id = instance.id(); + let entry = match self.instances.entry(id) { + Entry::Occupied(e) => return e.into_mut(), + Entry::Vacant(v) => v, + }; + + let cursor = Cursor::new(instance.wasm); + let module = decode_module(cursor).unwrap(); + #[cfg(watt_debug)] + crate::debug::print_module(&module); + let extern_vals = import::extern_vals(&module, &mut self.store); + let module_instance = + instantiate_module(&mut self.store, module, &extern_vals).unwrap(); + entry.insert(module_instance) + } + } + + STATE.with(|state| { + let mut state = state.borrow_mut(); + let instance = state.instance(instance); + let main = match get_export(instance, fun) { + Ok(ExternVal::Func(main)) => main, + _ => unimplemented!("unresolved macro: {:?}", fun), + }; + + let _guard = Data::guard(); + let args = Data::with(|d| { + inputs + .into_iter() + .map(|input| Value::I32(d.tokenstream.push(input))) + .collect() + }); + + let res = invoke_func(&mut state.store, main, args); + let values = match res { + Ok(values) => values, + Err(err) => panic!("{:?}", err), + }; + let handle = values.into_iter().next().unwrap(); + let handle = match handle { + Value::I32(handle) => handle, + _ => unimplemented!("unexpected macro return type"), + }; + Data::with(|d| d.tokenstream[handle].clone()) + }) } diff --git a/src/lib.rs b/src/lib.rs index bd7c0a1..97e0c24 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -198,119 +198,160 @@ mod sym; mod debug; use proc_macro::TokenStream; +use std::sync::atomic::{AtomicUsize, Ordering::SeqCst}; -/// A #\[proc_macro\] implemented in wasm! +/// An instantiation of a WebAssembly module used to invoke procedural macro +/// methods on the wasm module. /// -/// # Canonical macro implementation: /// -/// ``` -/// # const IGNORE: &str = stringify! { -/// use proc_macro2::TokenStream; -/// -/// #[no_mangle] -/// pub extern "C" fn my_macro(input: TokenStream) -> TokenStream { -/// proc_macro2::set_wasm_panic_hook(); -/// -/// ... -/// } -/// # }; -/// ``` -/// -/// # Canonical entry point: +/// # Examples /// /// ``` /// # const IGNORE: &str = stringify! { -/// extern crate proc_macro; -/// -/// use proc_macro::TokenStream; -/// -/// static WASM: &[u8] = include_bytes!("my_macro.wasm"); -/// -/// #[proc_macro] -/// pub fn my_macro(input: TokenStream) -> TokenStream { -/// watt::proc_macro("my_macro", input, WASM) -/// } +/// static WASM: watt::Instance = watt::Instance::new(include_bytes!("my_macro.wasm")); /// # }; /// ``` -pub fn proc_macro(fun: &str, input: TokenStream, wasm: &[u8]) -> TokenStream { - exec::proc_macro(fun, vec![input], wasm) +pub struct Instance { + wasm: &'static [u8], + id: AtomicUsize, } -/// A #\[proc_macro_derive\] implemented in wasm! -/// -/// # Canonical macro implementation: -/// -/// ``` -/// # const IGNORE: &str = stringify! { -/// use proc_macro2::TokenStream; -/// -/// #[no_mangle] -/// pub extern "C" fn my_macro(input: TokenStream) -> TokenStream { -/// proc_macro2::set_wasm_panic_hook(); -/// -/// ... -/// } -/// # }; -/// ``` -/// -/// # Canonical entry point: -/// -/// ``` -/// # const IGNORE: &str = stringify! { -/// extern crate proc_macro; -/// -/// use proc_macro::TokenStream; -/// -/// static WASM: &[u8] = include_bytes!("my_macro.wasm"); -/// -/// #[proc_macro_derive(MyDerive)] -/// pub fn my_macro(input: TokenStream) -> TokenStream { -/// watt::proc_macro_derive("my_macro", input, WASM) -/// } -/// # }; -/// ``` -pub fn proc_macro_derive(fun: &str, input: TokenStream, wasm: &[u8]) -> TokenStream { - exec::proc_macro(fun, vec![input], wasm) -} +impl Instance { + /// Creates a new `Instance` from the statically included blob of wasm bytes. + /// + /// # Examples + /// + /// ``` + /// # const IGNORE: &str = stringify! { + /// static WASM: watt::Instance = watt::Instance::new(include_bytes!("my_macro.wasm")); + /// # }; + /// ``` + pub const fn new(wasm: &'static [u8]) -> Instance { + Instance { + wasm, + id: AtomicUsize::new(0), + } + } -/// A #\[proc_macro_attribute\] implemented in wasm! -/// -/// # Canonical macro implementation: -/// -/// ``` -/// # const IGNORE: &str = stringify! { -/// use proc_macro2::TokenStream; -/// -/// #[no_mangle] -/// pub extern "C" fn my_macro(args: TokenStream, input: TokenStream) -> TokenStream { -/// proc_macro2::set_wasm_panic_hook(); -/// -/// ... -/// } -/// # }; -/// ``` -/// -/// # Canonical entry point: -/// -/// ``` -/// # const IGNORE: &str = stringify! { -/// extern crate proc_macro; -/// -/// use proc_macro::TokenStream; -/// -/// static WASM: &[u8] = include_bytes!("my_macro.wasm"); -/// -/// #[proc_macro_attribute] -/// pub fn my_macro(args: TokenStream, input: TokenStream) -> TokenStream { -/// watt::proc_macro_attribute("my_macro", args, input, WASM) -/// } -/// # }; -/// ``` -pub fn proc_macro_attribute( - fun: &str, - args: TokenStream, - input: TokenStream, - wasm: &[u8], -) -> TokenStream { - exec::proc_macro(fun, vec![args, input], wasm) + /// A #\[proc_macro\] implemented in wasm! + /// + /// # Canonical macro implementation: + /// + /// ``` + /// # const IGNORE: &str = stringify! { + /// use proc_macro2::TokenStream; + /// + /// #[proc_macro2::proc_macro] + /// pub fn my_macro(input: TokenStream) -> TokenStream { + /// ... + /// } + /// # }; + /// ``` + /// + /// # Canonical entry point: + /// + /// ``` + /// # const IGNORE: &str = stringify! { + /// extern crate proc_macro; + /// + /// use proc_macro::TokenStream; + /// + /// static WASM: watt::Instance = watt::Instance::new(include_bytes!("my_macro.wasm")); + /// + /// #[proc_macro] + /// pub fn my_macro(input: TokenStream) -> TokenStream { + /// WASM.proc_macro("my_macro", input) + /// } + /// # }; + /// ``` + pub fn proc_macro(&self, fun: &str, input: TokenStream) -> TokenStream { + exec::proc_macro(fun, vec![input], self) + } + + /// A #\[proc_macro_derive\] implemented in wasm! + /// + /// # Canonical macro implementation: + /// + /// ``` + /// # const IGNORE: &str = stringify! { + /// use proc_macro2::TokenStream; + /// + /// #[proc_macro2::proc_macro_derive(MyTrait)] + /// pub fn my_macro(input: TokenStream) -> TokenStream { + /// ... + /// } + /// # }; + /// ``` + /// + /// # Canonical entry point: + /// + /// ``` + /// # const IGNORE: &str = stringify! { + /// extern crate proc_macro; + /// + /// use proc_macro::TokenStream; + /// + /// static WASM: watt::Instance = watt::Instance::new(include_bytes!("my_macro.wasm")); + /// + /// #[proc_macro_derive(MyDerive)] + /// pub fn my_macro(input: TokenStream) -> TokenStream { + /// WASM.proc_macro_derive("my_macro", input) + /// } + /// # }; + /// ``` + pub fn proc_macro_derive(&self, fun: &str, input: TokenStream) -> TokenStream { + exec::proc_macro(fun, vec![input], self) + } + + /// A #\[proc_macro_attribute\] implemented in wasm! + /// + /// # Canonical macro implementation: + /// + /// ``` + /// # const IGNORE: &str = stringify! { + /// use proc_macro2::TokenStream; + /// + /// #[proc_macro2::proc_macro_attribute] + /// pub fn my_macro(args: TokenStream, input: TokenStream) -> TokenStream { + /// ... + /// } + /// # }; + /// ``` + /// + /// # Canonical entry point: + /// + /// ``` + /// # const IGNORE: &str = stringify! { + /// extern crate proc_macro; + /// + /// use proc_macro::TokenStream; + /// + /// static WASM: watt::Instance = watt::Instance::new(include_bytes!("my_macro.wasm")); + /// + /// #[proc_macro_attribute] + /// pub fn my_macro(args: TokenStream, input: TokenStream) -> TokenStream { + /// WASM.proc_macro_attribute("my_macro", args, input) + /// } + /// # }; + /// ``` + pub fn proc_macro_attribute( + &self, + fun: &str, + args: TokenStream, + input: TokenStream, + ) -> TokenStream { + exec::proc_macro(fun, vec![args, input], self) + } + + fn id(&self) -> usize { + static NEXT_ID: AtomicUsize = AtomicUsize::new(1); + match self.id.load(SeqCst) { + 0 => {} + n => return n, + } + let id = NEXT_ID.fetch_add(1, SeqCst); + self.id + .compare_exchange(0, id, SeqCst, SeqCst) + .unwrap_or_else(|id| id) + } }