feat: shared tables across modules

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

Author: explodingcamera

ARCHITECTURE.md

@@ -0,0 +1 @@
+# TinyWasm's Architecture

Cargo.lock

@@ -184,11 +184,6 @@ pub(crate) fn convert_blocktype(blocktype: wasmparser::BlockType) -> BlockArgs {
     use wasmparser::BlockType::*;
     match blocktype {
         Empty => BlockArgs::Empty,
-
-        // We should maybe have all this in a single variant for our custom bytecode
-
-        // TODO: maybe solve this differently so we can support 128-bit values
-        // without having to increase the size of the WasmValue enum
         Type(ty) => BlockArgs::Type(convert_valtype(&ty)),
         FuncType(ty) => BlockArgs::FuncType(ty),
     }
@@ -228,8 +223,8 @@ pub fn process_const_operator(op: wasmparser::Operator) -> Result<ConstInstructi
         wasmparser::Operator::RefFunc { function_index } => Ok(ConstInstruction::RefFunc(function_index)),
         wasmparser::Operator::I32Const { value } => Ok(ConstInstruction::I32Const(value)),
         wasmparser::Operator::I64Const { value } => Ok(ConstInstruction::I64Const(value)),
-        wasmparser::Operator::F32Const { value } => Ok(ConstInstruction::F32Const(f32::from_bits(value.bits()))), // TODO: check if this is correct
-        wasmparser::Operator::F64Const { value } => Ok(ConstInstruction::F64Const(f64::from_bits(value.bits()))), // TODO: check if this is correct
+        wasmparser::Operator::F32Const { value } => Ok(ConstInstruction::F32Const(f32::from_bits(value.bits()))),
+        wasmparser::Operator::F64Const { value } => Ok(ConstInstruction::F64Const(f64::from_bits(value.bits()))),
         wasmparser::Operator::GlobalGet { global_index } => Ok(ConstInstruction::GlobalGet(global_index)),
         op => Err(crate::ParseError::UnsupportedOperator(format!("Unsupported const instruction: {:?}", op))),
     }

crates/parser/src/conversion.rs

@@ -57,6 +57,7 @@ impl ModuleInstance {
         // don't need to create a auxiliary frame etc.
 
         let idx = store.next_module_instance_idx();
+        log::error!("Instantiating module at index {}", idx);
         let imports = imports.unwrap_or_default();
 
         let mut addrs = imports.link(store, &module, idx)?;
@@ -68,13 +69,12 @@ impl ModuleInstance {
         addrs.tables.extend(store.init_tables(data.table_types.into(), idx)?);
         addrs.memories.extend(store.init_memories(data.memory_types.into(), idx)?);
 
-        let elem_addrs = store.init_elements(&addrs.tables, data.elements.into(), idx)?;
+        let elem_addrs = store.init_elements(&addrs.tables, &addrs.funcs, data.elements.into(), idx)?;
         let data_addrs = store.init_datas(&addrs.memories, data.data.into(), idx)?;
 
         let instance = ModuleInstanceInner {
             store_id: store.id(),
             idx,
-
             types: data.func_types,
             func_addrs: addrs.funcs.into_boxed_slice(),
             table_addrs: addrs.tables.into_boxed_slice(),

crates/tinywasm/src/instance.rs

@@ -7,7 +7,7 @@ use crate::{
     CallFrame, Error, FuncContext, LabelArgs, ModuleInstance, Result, Store, Trap,
 };
 use alloc::{string::ToString, vec::Vec};
-use tinywasm_types::{ElementKind, Instruction};
+use tinywasm_types::{ElementKind, Instruction, ValType};
 
 mod macros;
 mod traits;
@@ -171,15 +171,17 @@ fn exec_one(
             let call_ty = module.func_ty(*type_addr);
 
             let func_idx = stack.values.pop_t::<u32>()?;
-            let actual_func_addr = table
+
+            // verify that the table is of the right type, this should be validated by the parser already
+            assert!(table.borrow().kind.element_type == ValType::RefFunc, "table is not of type funcref");
+
+            let func_ref = table
                 .borrow()
                 .get(func_idx as usize)?
+                .addr()
                 .ok_or_else(|| Trap::UninitializedElement { index: func_idx as usize })?;
 
-            let resolved_func_addr = module.resolve_func_addr(actual_func_addr);
-
-            // prepare the call frame
-            let func_inst = store.get_func(resolved_func_addr as usize)?;
+            let func_inst = store.get_func(func_ref as usize)?;
             let func_ty = func_inst.func.ty();
 
             let func = match &func_inst.func {
@@ -200,8 +202,7 @@ fn exec_one(
             }
 
             let params = stack.values.pop_n_rev(func_ty.params.len())?;
-            let call_frame =
-                CallFrame::new_raw(resolved_func_addr as usize, &params, func.locals.to_vec(), func_inst._owner);
+            let call_frame = CallFrame::new_raw(func_ref as usize, &params, func.locals.to_vec(), func_inst._owner);
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
@@ -606,7 +607,7 @@ fn exec_one(
                 return Err(Trap::TableOutOfBounds { offset: 0, len: 0, max: 0 }.into());
             };
 
-            table.borrow_mut().init(0, items)?;
+            table.borrow_mut().init(module.func_addrs(), 0, items)?;
         }
 
         I32TruncSatF32S => arithmetic_single!(trunc, f32, i32, stack),

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

@@ -82,7 +82,6 @@ impl Default for Store {
 ///
 /// Data should only be addressable by the module that owns it
 /// See <https://webassembly.github.io/spec/core/exec/runtime.html#store>
-// TODO: Arena allocate these?
 pub(crate) struct StoreData {
     pub(crate) funcs: Vec<Rc<FunctionInstance>>,
     pub(crate) tables: Vec<Rc<RefCell<TableInstance>>>,
@@ -178,6 +177,7 @@ impl Store {
     pub(crate) fn init_elements(
         &mut self,
         table_addrs: &[TableAddr],
+        func_addrs: &[FuncAddr],
         elements: Vec<Element>,
         idx: ModuleInstanceAddr,
     ) -> Result<Box<[Addr]>> {
@@ -194,6 +194,8 @@ impl Store {
                 })
                 .collect::<Result<Vec<_>>>()?;
 
+            log::error!("element kind: {:?}", element.kind);
+
             let items = match element.kind {
                 // doesn't need to be initialized, can be initialized lazily using the `table.init` instruction
                 ElementKind::Passive => Some(init),
@@ -218,7 +220,7 @@ impl Store {
                     // d. Execute the instruction i32.const n
                     // e. Execute the instruction table.init tableidx i
                     if let Some(table) = self.data.tables.get_mut(table_addr as usize) {
-                        table.borrow_mut().init(offset, &init)?;
+                        table.borrow_mut().init(func_addrs, offset, &init)?;
                     } else {
                         log::error!("table {} not found", table);
                     }
@@ -385,8 +387,6 @@ impl Store {
 /// See <https://webassembly.github.io/spec/core/exec/runtime.html#function-instances>
 pub struct FunctionInstance {
     pub(crate) func: Function,
-
-    // TODO: this is important for call_indirect
     pub(crate) _type_idx: TypeAddr,
     pub(crate) _owner: ModuleInstanceAddr, // index into store.module_instances, none for host functions
 }
@@ -408,52 +408,77 @@ impl FunctionInstance {
     }
 }
 
+#[derive(Debug, Clone, Copy)]
+pub(crate) enum TableElement {
+    Uninitialized,
+    Initialized(Addr),
+}
+
+impl TableElement {
+    pub(crate) fn addr(&self) -> Option<Addr> {
+        match self {
+            TableElement::Uninitialized => None,
+            TableElement::Initialized(addr) => Some(*addr),
+        }
+    }
+}
+
 /// A WebAssembly Table Instance
 ///
 /// See <https://webassembly.github.io/spec/core/exec/runtime.html#table-instances>
 #[derive(Debug)]
 pub(crate) struct TableInstance {
-    pub(crate) elements: Vec<Option<Addr>>,
+    pub(crate) elements: Vec<TableElement>,
     pub(crate) kind: TableType,
     pub(crate) _owner: ModuleInstanceAddr, // index into store.module_instances
 }
 
 impl TableInstance {
     pub(crate) fn new(kind: TableType, owner: ModuleInstanceAddr) -> Self {
-        Self { elements: vec![None; kind.size_initial as usize], kind, _owner: owner }
+        Self { elements: vec![TableElement::Uninitialized; kind.size_initial as usize], kind, _owner: owner }
     }
 
     pub(crate) fn get_wasm_val(&self, addr: usize) -> Result<WasmValue> {
+        let val = self.get(addr)?.addr();
+
         Ok(match self.kind.element_type {
-            ValType::RefFunc => {
-                self.get(addr)?.map(|v| WasmValue::RefFunc(v)).unwrap_or(WasmValue::RefNull(ValType::RefFunc))
-            }
+            ValType::RefFunc => val.map(|v| WasmValue::RefFunc(v)).unwrap_or(WasmValue::RefNull(ValType::RefFunc)),
             ValType::RefExtern => {
-                self.get(addr)?.map(|v| WasmValue::RefExtern(v)).unwrap_or(WasmValue::RefNull(ValType::RefExtern))
+                val.map(|v| WasmValue::RefExtern(v)).unwrap_or(WasmValue::RefNull(ValType::RefExtern))
             }
             _ => unimplemented!("unsupported table type: {:?}", self.kind.element_type),
         })
     }
 
-    pub(crate) fn get(&self, addr: usize) -> Result<Option<Addr>> {
-        self.elements.get(addr).copied().ok_or_else(|| Error::Trap(Trap::UndefinedElement { index: addr }))
+    pub(crate) fn get(&self, addr: usize) -> Result<&TableElement> {
+        self.elements.get(addr).ok_or_else(|| Error::Trap(Trap::UndefinedElement { index: addr }))
     }
 
-    pub(crate) fn set(&mut self, addr: usize, value: Addr) -> Result<()> {
-        if addr >= self.elements.len() {
-            return Err(Error::Other(format!("table element {} not found", addr)));
-        }
-
-        self.elements[addr] = Some(value);
+    pub(crate) fn set(&mut self, table_idx: usize, value: Addr) -> Result<()> {
+        let el = self
+            .elements
+            .get_mut(table_idx)
+            .ok_or_else(|| Error::Other(format!("table element {} not found", table_idx)))?;
+        *el = TableElement::Initialized(value);
         Ok(())
     }
 
     pub(crate) fn size(&self) -> i32 {
         self.elements.len() as i32
     }
 
-    pub(crate) fn init(&mut self, offset: i32, init: &[Addr]) -> Result<()> {
-        let init = init.iter().map(|item| Some(*item)).collect::<Vec<_>>();
+    pub(crate) fn init(&mut self, func_addrs: &[u32], offset: i32, init: &[Addr]) -> Result<()> {
+        let init = init
+            .iter()
+            .map(|item| {
+                TableElement::Initialized(match self.kind.element_type == ValType::RefFunc {
+                    true => *func_addrs.get(*item as usize).expect(
+                        "error initializing table: function not found. This should have been caught by the validator",
+                    ),
+                    false => *item,
+                })
+            })
+            .collect::<Vec<_>>();
 
         let offset = offset as usize;
         let end = offset.checked_add(init.len()).ok_or_else(|| {
@@ -465,6 +490,7 @@ impl TableInstance {
         }
 
         self.elements[offset..end].copy_from_slice(&init);
+        log::debug!("table: {:?}", self.elements);
         Ok(())
     }
 }