chore: refactor executer, change the module when executing imported functions

Signed-off-by: Henry Gressmann <mail@henrygressmann.de>

Author: explodingcamera

crates/tinywasm/src/func.rs

@@ -27,7 +27,7 @@ impl FuncHandle {
 
         // 1. Assert: funcs[func_addr] exists
         // 2. let func_inst be the functiuon instance funcs[func_addr]
-        let func_inst = store.get_func(self.addr as usize)?;
+        let func_inst = store.get_func(self.addr as usize)?.clone();
 
         // 3. Let func_ty be the function type
         let func_ty = &self.ty;
@@ -52,26 +52,26 @@ impl FuncHandle {
             }
         }
 
-        let wasm_func = match &func_inst.func {
+        let locals = match &func_inst.func {
             crate::Function::Host(h) => {
                 let func = h.func.clone();
                 let ctx = FuncContext { store, module: &self.module };
                 return (func)(ctx, params);
             }
-            crate::Function::Wasm(ref f) => f,
+            crate::Function::Wasm(ref f) => f.locals.to_vec(),
         };
 
         // 6. Let f be the dummy frame
-        debug!("locals: {:?}", wasm_func.locals);
-        let call_frame = CallFrame::new(self.addr as usize, params, wasm_func.locals.to_vec(), self.module.id());
+        debug!("locals: {:?}", locals);
+        let call_frame = CallFrame::new(func_inst, params, locals);
 
         // 7. Push the frame f to the call stack
         // & 8. Push the values to the stack (Not needed since the call frame owns the values)
         stack.call_stack.push(call_frame)?;
 
         // 9. Invoke the function instance
         let runtime = store.runtime();
-        runtime.exec(store, &mut stack, self.module.clone())?;
+        runtime.exec(store, &mut stack)?;
 
         // Once the function returns:
         let result_m = func_ty.results.len();

crates/tinywasm/src/instance.rs

@@ -81,7 +81,7 @@ impl ModuleInstance {
             mem_addrs: addrs.memories.into_boxed_slice(),
             global_addrs: addrs.globals.into_boxed_slice(),
             elem_addrs,
-            data_addrs: data_addrs,
+            data_addrs,
             func_start: data.start_func,
             imports: data.imports,
             exports: data.exports,
@@ -111,11 +111,7 @@ impl ModuleInstance {
         &self.0.func_addrs
     }
 
-    pub(crate) fn _global_addrs(&self) -> &[GlobalAddr] {
-        &self.0.global_addrs
-    }
-
-    pub(crate) fn func_ty_addrs(&self) -> &[FuncType] {
+    pub(crate) fn func_tys(&self) -> &[FuncType] {
         &self.0.types
     }
 

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

@@ -15,32 +15,32 @@ use macros::*;
 use traits::*;
 
 impl DefaultRuntime {
-    pub(crate) fn exec(&self, store: &mut Store, stack: &mut Stack, module: ModuleInstance) -> Result<()> {
-        log::debug!("func_addrs: {:?}", module.func_addrs());
-        log::debug!("func_ty_addrs: {:?}", module.func_ty_addrs().len());
-        log::debug!("store funcs: {:?}", store.data.funcs.len());
-
+    pub(crate) fn exec(&self, store: &mut Store, stack: &mut Stack) -> Result<()> {
         // The current call frame, gets updated inside of exec_one
         let mut cf = stack.call_stack.pop()?;
 
-        // The function to execute, gets updated from ExecResult::Call
-        let mut func_inst = store.get_func(cf.func_ptr)?.clone();
+        let mut func_inst = cf.func_instance.clone();
         let mut wasm_func = func_inst.assert_wasm().expect("exec expected wasm function");
+
+        // The function to execute, gets updated from ExecResult::Call
         let mut instrs = &wasm_func.instructions;
 
-        let mut current_module = module;
+        let mut current_module = store
+            .get_module_instance(func_inst.owner)
+            .expect("exec expected module instance to exist for function")
+            .clone();
 
         while let Some(instr) = instrs.get(cf.instr_ptr) {
             match exec_one(&mut cf, instr, instrs, stack, store, &current_module)? {
                 // Continue execution at the new top of the call stack
                 ExecResult::Call => {
                     cf = stack.call_stack.pop()?;
-                    func_inst = store.get_func(cf.func_ptr)?.clone();
-                    wasm_func = func_inst.assert_wasm().expect("call expected wasm function");
+                    func_inst = cf.func_instance.clone();
+                    wasm_func = func_inst.assert_wasm().expect("exec expected wasm function");
                     instrs = &wasm_func.instructions;
 
-                    if cf.module != current_module.id() {
-                        current_module.swap(store.get_module_instance(cf.module).unwrap().clone());
+                    if cf.func_instance.owner != current_module.id() {
+                        current_module.swap(store.get_module_instance(cf.func_instance.owner).unwrap().clone());
                     }
 
                     continue;
@@ -135,10 +135,10 @@ fn exec_one(
             log::info!("start call");
             // prepare the call frame
             let func_idx = module.resolve_func_addr(*v);
-            let func_inst = store.get_func(func_idx as usize)?;
+            let func_inst = store.get_func(func_idx as usize)?.clone();
 
-            let func = match &func_inst.func {
-                crate::Function::Wasm(ref f) => f,
+            let (locals, ty) = match &func_inst.func {
+                crate::Function::Wasm(ref f) => (f.locals.to_vec(), f.ty.clone()),
                 crate::Function::Host(host_func) => {
                     let func = host_func.func.clone();
                     log::info!("Getting params: {:?}", host_func.ty.params);
@@ -150,8 +150,11 @@ fn exec_one(
                 }
             };
 
-            let params = stack.values.pop_n_rev(func.ty.params.len())?;
-            let call_frame = CallFrame::new_raw(func_idx as usize, &params, func.locals.to_vec(), func_inst._owner);
+            let params = stack.values.pop_n_rev(ty.params.len())?;
+            log::info!("call: current fn owner: {:?}", module.id());
+            log::info!("call: func owner: {:?}", func_inst.owner);
+
+            let call_frame = CallFrame::new_raw(func_inst, &params, locals);
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
@@ -163,29 +166,39 @@ fn exec_one(
         }
 
         CallIndirect(type_addr, table_addr) => {
-            let table_idx = module.resolve_table_addr(*table_addr);
-            let table = store.get_table(table_idx as usize)?;
-
-            // TODO: currently, the type resolution is subtlely broken for imported functions
-
-            let call_ty = module.func_ty(*type_addr);
-
-            let func_idx = stack.values.pop_t::<u32>()?;
+            let table = store.get_table(module.resolve_table_addr(*table_addr) as usize)?;
+            let table_idx = stack.values.pop_t::<u32>()?;
 
             // 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 func_ref = {
+                table
+                    .borrow()
+                    .get(table_idx as usize)?
+                    .addr()
+                    .ok_or(Trap::UninitializedElement { index: table_idx as usize })?
+            };
 
-            let func_inst = store.get_func(func_ref as usize)?;
+            let func_inst = store.get_func(func_ref as usize)?.clone();
             let func_ty = func_inst.func.ty();
 
-            let func = match &func_inst.func {
-                crate::Function::Wasm(ref f) => f,
+            log::info!("type_addr: {}", type_addr);
+            log::info!("types: {:?}", module.func_tys());
+            let call_ty = module.func_ty(*type_addr);
+
+            log::info!("call_indirect: current fn owner: {:?}", module.id());
+            log::info!("call_indirect: func owner: {:?}", func_inst.owner);
+
+            if func_ty != call_ty {
+                log::error!("indirect call type mismatch: {:?} != {:?}", func_ty, call_ty);
+                return Err(
+                    Trap::IndirectCallTypeMismatch { actual: func_ty.clone(), expected: call_ty.clone() }.into()
+                );
+            }
+
+            let locals = match &func_inst.func {
+                crate::Function::Wasm(ref f) => f.locals.to_vec(),
                 crate::Function::Host(host_func) => {
                     let func = host_func.func.clone();
                     let params = stack.values.pop_params(&func_ty.params)?;
@@ -195,14 +208,8 @@ fn exec_one(
                 }
             };
 
-            if func_ty != call_ty {
-                return Err(
-                    Trap::IndirectCallTypeMismatch { actual: func_ty.clone(), expected: call_ty.clone() }.into()
-                );
-            }
-
             let params = stack.values.pop_n_rev(func_ty.params.len())?;
-            let call_frame = CallFrame::new_raw(func_ref as usize, &params, func.locals.to_vec(), func_inst._owner);
+            let call_frame = CallFrame::new_raw(func_inst, &params, locals);
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction

crates/tinywasm/src/runtime/stack/call_stack.rs

@@ -1,6 +1,6 @@
-use crate::{runtime::RawWasmValue, BlockType, Error, LabelFrame, Result, Trap};
-use alloc::{boxed::Box, vec::Vec};
-use tinywasm_types::{ModuleInstanceAddr, ValType, WasmValue};
+use crate::{runtime::RawWasmValue, BlockType, Error, FunctionInstance, LabelFrame, Result, Trap};
+use alloc::{boxed::Box, rc::Rc, vec::Vec};
+use tinywasm_types::{ValType, WasmValue};
 
 use super::blocks::Labels;
 
@@ -52,9 +52,9 @@ impl CallStack {
 
 #[derive(Debug, Clone)]
 pub(crate) struct CallFrame {
-    pub(crate) module: ModuleInstanceAddr,
     pub(crate) instr_ptr: usize,
-    pub(crate) func_ptr: usize,
+    // pub(crate) module: ModuleInstanceAddr,
+    pub(crate) func_instance: Rc<FunctionInstance>,
 
     pub(crate) labels: Labels,
     pub(crate) locals: Box<[RawWasmValue]>,
@@ -109,36 +109,32 @@ impl CallFrame {
     }
 
     pub(crate) fn new_raw(
-        func_ptr: usize,
+        func_instance_ptr: Rc<FunctionInstance>,
         params: &[RawWasmValue],
         local_types: Vec<ValType>,
-        module: ModuleInstanceAddr,
     ) -> Self {
         let mut locals = Vec::with_capacity(local_types.len() + params.len());
         locals.extend(params.iter().cloned());
         locals.extend(local_types.iter().map(|_| RawWasmValue::default()));
 
         Self {
             instr_ptr: 0,
-            func_ptr,
+            func_instance: func_instance_ptr,
             local_count: locals.len(),
             locals: locals.into_boxed_slice(),
             labels: Labels::default(),
-            module,
         }
     }
 
     pub(crate) fn new(
-        func_ptr: usize,
+        func_instance_ptr: Rc<FunctionInstance>,
         params: &[WasmValue],
         local_types: Vec<ValType>,
-        module: ModuleInstanceAddr,
     ) -> Self {
         CallFrame::new_raw(
-            func_ptr,
+            func_instance_ptr,
             &params.iter().map(|v| RawWasmValue::from(*v)).collect::<Vec<_>>(),
             local_types,
-            module,
         )
     }
 

crates/tinywasm/src/store.rs

@@ -121,7 +121,7 @@ impl Store {
             self.data.funcs.push(Rc::new(FunctionInstance {
                 func: Function::Wasm(func),
                 _type_idx: type_idx,
-                _owner: idx,
+                owner: idx,
             }));
             func_addrs.push((i + func_count) as FuncAddr);
         }
@@ -303,7 +303,7 @@ impl Store {
     }
 
     pub(crate) fn add_func(&mut self, func: Function, type_idx: TypeAddr, idx: ModuleInstanceAddr) -> Result<FuncAddr> {
-        self.data.funcs.push(Rc::new(FunctionInstance { func, _type_idx: type_idx, _owner: idx }));
+        self.data.funcs.push(Rc::new(FunctionInstance { func, _type_idx: type_idx, owner: idx }));
         Ok(self.data.funcs.len() as FuncAddr - 1)
     }
 
@@ -388,7 +388,7 @@ impl Store {
 pub struct FunctionInstance {
     pub(crate) func: Function,
     pub(crate) _type_idx: TypeAddr,
-    pub(crate) _owner: ModuleInstanceAddr, // index into store.module_instances, none for host functions
+    pub(crate) owner: ModuleInstanceAddr, // index into store.module_instances, none for host functions
 }
 
 // TODO: check if this actually helps
@@ -442,10 +442,8 @@ impl TableInstance {
         let val = self.get(addr)?.addr();
 
         Ok(match self.kind.element_type {
-            ValType::RefFunc => val.map(|v| WasmValue::RefFunc(v)).unwrap_or(WasmValue::RefNull(ValType::RefFunc)),
-            ValType::RefExtern => {
-                val.map(|v| WasmValue::RefExtern(v)).unwrap_or(WasmValue::RefNull(ValType::RefExtern))
-            }
+            ValType::RefFunc => val.map(WasmValue::RefFunc).unwrap_or(WasmValue::RefNull(ValType::RefFunc)),
+            ValType::RefExtern => val.map(WasmValue::RefExtern).unwrap_or(WasmValue::RefNull(ValType::RefExtern)),
             _ => unimplemented!("unsupported table type: {:?}", self.kind.element_type),
         })
     }