chore: switch module instance when switching call stack

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

crates/tinywasm/src/func.rs

@@ -63,7 +63,7 @@ impl FuncHandle {
 
         // 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());
+        let call_frame = CallFrame::new(self.addr as usize, params, wasm_func.locals.to_vec(), self.module.id());
 
         // 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)

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

@@ -28,16 +28,23 @@ impl DefaultRuntime {
         let mut wasm_func = func_inst.assert_wasm().expect("exec expected wasm function");
         let mut instrs = &wasm_func.instructions;
 
+        let mut current_module = module;
+
         // TODO: we might be able to index into the instructions directly
         // since the instruction pointer should always be in bounds
         while let Some(instr) = instrs.get(cf.instr_ptr) {
-            match exec_one(&mut cf, instr, instrs, stack, store, &module)? {
+            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");
                     instrs = &wasm_func.instructions;
+
+                    if cf.module != current_module.id() {
+                        current_module = store.get_module_instance(cf.module).unwrap().clone()
+                    }
+
                     continue;
                 }
 
@@ -146,7 +153,7 @@ 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());
+            let call_frame = CallFrame::new_raw(func_idx as usize, &params, func.locals.to_vec(), func_inst._owner);
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
@@ -189,7 +196,8 @@ 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());
+            let call_frame =
+                CallFrame::new_raw(resolved_func_addr as usize, &params, func.locals.to_vec(), func_inst._owner);
 
             // 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::{ValType, WasmValue};
+use tinywasm_types::{ModuleInstanceAddr, ValType, WasmValue};
 
 use super::blocks::Labels;
 
@@ -52,6 +52,7 @@ impl CallStack {
 
 #[derive(Debug, Clone)]
 pub(crate) struct CallFrame {
+    pub(crate) module: ModuleInstanceAddr,
     pub(crate) instr_ptr: usize,
     pub(crate) func_ptr: usize,
 
@@ -107,7 +108,12 @@ impl CallFrame {
         Some(())
     }
 
-    pub(crate) fn new_raw(func_ptr: usize, params: &[RawWasmValue], local_types: Vec<ValType>) -> Self {
+    pub(crate) fn new_raw(
+        func_ptr: usize,
+        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()));
@@ -118,11 +124,22 @@ impl CallFrame {
             local_count: locals.len(),
             locals: locals.into_boxed_slice(),
             labels: Labels::default(),
+            module,
         }
     }
 
-    pub(crate) fn new(func_ptr: usize, params: &[WasmValue], local_types: Vec<ValType>) -> Self {
-        CallFrame::new_raw(func_ptr, &params.iter().map(|v| RawWasmValue::from(*v)).collect::<Vec<_>>(), local_types)
+    pub(crate) fn new(
+        func_ptr: usize,
+        params: &[WasmValue],
+        local_types: Vec<ValType>,
+        module: ModuleInstanceAddr,
+    ) -> Self {
+        CallFrame::new_raw(
+            func_ptr,
+            &params.iter().map(|v| RawWasmValue::from(*v)).collect::<Vec<_>>(),
+            local_types,
+            module,
+        )
     }
 
     #[inline]

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

@@ -10,8 +10,6 @@ pub(crate) const STACK_SIZE: usize = 1024;
 #[derive(Debug)]
 pub(crate) struct ValueStack {
     stack: Vec<RawWasmValue>,
-
-    // TODO: don't pop the stack, just keep track of the top for better performance
     top: usize,
 }