chore: progress towards linking modules

Signed-off-by: Henry Gressmann <[email protected]>

crates/tinywasm/src/imports.rs

@@ -3,14 +3,16 @@ use alloc::{
     collections::BTreeMap,
     string::{String, ToString},
 };
-use tinywasm_types::{Global, GlobalType, WasmValue};
+use tinywasm_types::{Global, GlobalType, ModuleInstanceAddr, WasmValue};
 
 #[derive(Debug)]
 #[non_exhaustive]
 /// An external value
 pub enum Extern {
     /// A global value
     Global(Global),
+    /// A registered module
+    Module(String),
     // Func(HostFunc),
     // Table(Table),
 }
@@ -39,16 +41,39 @@ pub struct ExternName {
 /// Imports for a module instance
 pub struct Imports {
     values: BTreeMap<ExternName, Extern>,
+    modules: BTreeMap<String, ModuleInstanceAddr>,
+}
+
+pub(crate) struct LinkedImports {
+    pub(crate) values: BTreeMap<ExternName, Extern>,
+}
+
+impl LinkedImports {
+    pub(crate) fn get(&self, module: &str, name: &str) -> Option<&Extern> {
+        self.values.get(&ExternName {
+            module: module.to_string(),
+            name: name.to_string(),
+        })
+    }
 }
 
 impl Imports {
     /// Create a new empty import set
     pub fn new() -> Self {
         Imports {
             values: BTreeMap::new(),
+            modules: BTreeMap::new(),
         }
     }
 
+    /// Link a module
+    ///
+    /// This will automatically link all imported values
+    pub fn link_module(&mut self, name: &str, addr: ModuleInstanceAddr) -> Result<&mut Self> {
+        self.modules.insert(name.to_string(), addr);
+        Ok(self)
+    }
+
     /// Define an import
     pub fn define(&mut self, module: &str, name: &str, value: Extern) -> Result<&mut Self> {
         self.values.insert(
@@ -61,10 +86,8 @@ impl Imports {
         Ok(self)
     }
 
-    pub(crate) fn get(&self, module: &str, name: &str) -> Option<&Extern> {
-        self.values.get(&ExternName {
-            module: module.to_string(),
-            name: name.to_string(),
-        })
+    pub(crate) fn link(self, store: &mut crate::Store, module: &crate::Module) -> Result<LinkedImports> {
+        let values = self.values;
+        Ok(LinkedImports { values })
     }
 }

crates/tinywasm/src/instance.rs

@@ -36,17 +36,22 @@ pub(crate) struct ModuleInstanceInner {
 }
 
 impl ModuleInstance {
+    /// Get the module instance's address
+    pub fn id(&self) -> ModuleInstanceAddr {
+        self.0.idx
+    }
+
     /// Instantiate the module in the given store
     pub fn instantiate(store: &mut Store, module: Module, imports: Option<Imports>) -> Result<Self> {
         let idx = store.next_module_instance_idx();
         let imports = imports.unwrap_or_default();
+        let linked_imports = imports.link(store, &module)?;
 
         let func_addrs = store.add_funcs(module.data.funcs.into(), idx);
         let table_addrs = store.add_tables(module.data.table_types.into(), idx);
         let mem_addrs = store.add_mems(module.data.memory_types.into(), idx)?;
-
-        let global_addrs = store.add_globals(module.data.globals.into(), &module.data.imports, &imports, idx)?;
-        let elem_addrs = store.add_elems(module.data.elements.into(), idx);
+        let global_addrs = store.add_globals(module.data.globals.into(), &module.data.imports, &linked_imports, idx)?;
+        let elem_addrs = store.add_elems(module.data.elements.into(), idx)?;
         let data_addrs = store.add_datas(module.data.data.into(), idx);
 
         let instance = ModuleInstanceInner {

crates/tinywasm/src/runtime/executor/macros.rs

@@ -1,6 +1,8 @@
 //! More generic macros for various instructions
 //!
 //! These macros are used to generate the actual instruction implementations.
+//! In some basic tests this generated better assembly than using generic functions, even when inlined.
+//! (Something to revisit in the future)
 
 /// Load a value from memory
 macro_rules! mem_load {
@@ -37,6 +39,7 @@ macro_rules! mem_store {
     }};
 
     ($store_type:ty, $target_type:ty, $arg:ident, $stack:ident, $store:ident, $module:ident) => {{
+        // likewise, there could be a lot of performance improvements here
         let mem_idx = $module.resolve_mem_addr($arg.mem_addr);
         let mem = $store.get_mem(mem_idx as usize)?;
 

crates/tinywasm/src/runtime/executor/mod.rs

@@ -164,10 +164,6 @@ fn exec_one(
             }
         }
 
-        // Else(_end_offset) => {
-        //     // end the if block
-        //     cf.break_to(0, &mut stack.values)?;
-        // }
         Loop(args, end_offset) => {
             // let params = stack.values.pop_block_params(*args, &module)?;
             cf.enter_label(
@@ -240,6 +236,7 @@ fn exec_one(
             }
         }
 
+        // We're essentially using else as a EndBlockFrame instruction
         Else(end_offset) => {
             let Some(block) = cf.labels.pop() else {
                 panic!("else: no label to end, this should have been validated by the parser");

crates/tinywasm/src/runtime/stack.rs

@@ -9,10 +9,6 @@ pub(crate) use call_stack::CallFrame;
 /// A WebAssembly Stack
 #[derive(Debug, Default)]
 pub struct Stack {
-    // keeping this typed for now to make it easier to debug
-    // TODO: Maybe split into Vec<u8> and Vec<ValType> for better memory usage?
     pub(crate) values: ValueStack,
-
-    /// The call stack
     pub(crate) call_stack: CallStack,
 }

crates/tinywasm/src/store.rs

@@ -13,7 +13,7 @@ use tinywasm_types::{
 
 use crate::{
     runtime::{self, DefaultRuntime},
-    Error, Extern, Imports, ModuleInstance, RawWasmValue, Result,
+    Error, Extern, LinkedImports, ModuleInstance, RawWasmValue, Result,
 };
 
 // global store id counter
@@ -153,7 +153,7 @@ impl Store {
         &mut self,
         globals: Vec<Global>,
         wasm_imports: &[Import],
-        user_imports: &Imports,
+        user_imports: &LinkedImports,
         idx: ModuleInstanceAddr,
     ) -> Result<Vec<Addr>> {
         // TODO: initialize imported globals
@@ -191,41 +191,58 @@ impl Store {
         let iterator = imported_globals.into_iter().chain(globals.as_ref());
 
         for (i, global) in iterator.enumerate() {
-            use tinywasm_types::ConstInstruction::*;
-            let val = match global.init {
-                F32Const(f) => RawWasmValue::from(f),
-                F64Const(f) => RawWasmValue::from(f),
-                I32Const(i) => RawWasmValue::from(i),
-                I64Const(i) => RawWasmValue::from(i),
-                GlobalGet(addr) => {
-                    let addr = global_addrs[addr as usize];
-                    let global = self.data.globals[addr as usize].clone();
-                    let val = global.borrow().value;
-                    val
-                }
-                RefNull(_) => RawWasmValue::default(),
-                RefFunc(idx) => RawWasmValue::from(idx as i64),
-            };
-
-            self.data
-                .globals
-                .push(Rc::new(RefCell::new(GlobalInstance::new(global.ty, val, idx))));
-
+            self.data.globals.push(Rc::new(RefCell::new(GlobalInstance::new(
+                global.ty,
+                self.eval_const(&global.init)?,
+                idx,
+            ))));
             global_addrs.push((i + global_count) as Addr);
         }
-        log::debug!("global_addrs: {:?}", global_addrs);
+
         Ok(global_addrs)
     }
 
+    pub(crate) fn eval_const(&self, const_instr: &tinywasm_types::ConstInstruction) -> Result<RawWasmValue> {
+        use tinywasm_types::ConstInstruction::*;
+        let val = match const_instr {
+            F32Const(f) => RawWasmValue::from(*f),
+            F64Const(f) => RawWasmValue::from(*f),
+            I32Const(i) => RawWasmValue::from(*i),
+            I64Const(i) => RawWasmValue::from(*i),
+            GlobalGet(addr) => {
+                let addr = *addr as usize;
+                let global = self.data.globals[addr].clone();
+                let val = global.borrow().value;
+                val
+            }
+            RefNull(_) => RawWasmValue::default(),
+            RefFunc(idx) => RawWasmValue::from(*idx as i64),
+        };
+        Ok(val)
+    }
+
     /// Add elements to the store, returning their addresses in the store
-    pub(crate) fn add_elems(&mut self, elems: Vec<Element>, idx: ModuleInstanceAddr) -> Vec<Addr> {
+    /// Should be called after the tables have been added
+    pub(crate) fn add_elems(&mut self, elems: Vec<Element>, idx: ModuleInstanceAddr) -> Result<Vec<Addr>> {
         let elem_count = self.data.elems.len();
         let mut elem_addrs = Vec::with_capacity(elem_count);
         for (i, elem) in elems.into_iter().enumerate() {
             self.data.elems.push(ElemInstance::new(elem.kind, idx));
             elem_addrs.push((i + elem_count) as Addr);
+
+            // match elem.kind {
+            //     ElementKind::Active { table, offset } => {
+            //         // let table = self.data.tables[table as usize];
+
+            //         // let offset = self.eval_const(&offset)?;
+            //         // let offset = offset.raw_value() as usize;
+            //         // let offset = offset + elem_addrs[i] as usize;
+            //         // let offset = offset as Addr;
+            //     }
+            // }
         }
-        elem_addrs
+
+        Ok(elem_addrs)
     }
 
     /// Add data to the store, returning their addresses in the store
@@ -313,8 +330,8 @@ pub(crate) struct TableInstance {
 impl TableInstance {
     pub(crate) fn new(kind: TableType, owner: ModuleInstanceAddr) -> Self {
         Self {
+            elements: vec![0; kind.size_initial as usize],
             kind,
-            elements: Vec::new(),
             owner,
         }
     }