simple LUT linker

src/execution/execution_info.rs

@@ -9,11 +9,11 @@ pub struct ExecutionInfo<'r> {
     pub wasm_bytecode: &'r [u8],
     pub wasm_reader: WasmReader<'r>,
     pub fn_types: &'r [FuncType],
-    pub store: &'r mut Store,
+    pub store: Store,
 }
 
 impl<'r> ExecutionInfo<'r> {
-    pub fn new(wasm_bytecode: &'r [u8], fn_types: &'r [FuncType], store: &'r mut Store) -> Self {
+    pub fn new(wasm_bytecode: &'r [u8], fn_types: &'r [FuncType], store: Store) -> Self {
         ExecutionInfo {
             wasm_bytecode,
             wasm_reader: WasmReader::new(wasm_bytecode),

src/execution/interpreter_loop.rs

@@ -22,6 +22,7 @@ use crate::{
         },
     },
     execution::execution_info::ExecutionInfo,
+    execution::LUT,
     locals::Locals,
     store::{FuncInst, Store},
     value,
@@ -36,6 +37,7 @@ use crate::execution::hooks::HookSet;
 pub(super) fn run<H: HookSet>(
     modules: &mut [ExecutionInfo],
     current_module: &mut usize,
+    lut: &LUT,
     stack: &mut Stack,
     mut hooks: H,
 ) -> Result<(), RuntimeError> {
@@ -63,7 +65,8 @@ pub(super) fn run<H: HookSet>(
 
         match first_instr_byte {
             END => {
-                let maybe_return_address = stack.pop_stackframe();
+                let (return_module, maybe_return_address) = stack.pop_stackframe();
+                *current_module = return_module;
 
                 // We finished this entire invocation if there is no stackframe left. If there are
                 // one or more stack frames, we need to continue from where the callee was called
@@ -73,6 +76,7 @@ pub(super) fn run<H: HookSet>(
                 }
 
                 trace!("end of function reached, returning to previous stack frame");
+                wasm = &mut modules[return_module].wasm_reader;
                 wasm.pc = maybe_return_address;
             }
             RETURN => {
@@ -104,7 +108,10 @@ pub(super) fn run<H: HookSet>(
                 }
 
                 trace!("end of function reached, returning to previous stack frame");
-                wasm.pc = stack.pop_stackframe();
+                let (return_module, return_pc) = stack.pop_stackframe();
+                *current_module = return_module;
+                wasm = &mut modules[return_module].wasm_reader;
+                wasm.pc = return_pc;
             }
             CALL => {
                 let func_to_call_idx = wasm.read_var_u32().unwrap_validated() as FuncIdx;
@@ -128,12 +135,43 @@ pub(super) fn run<H: HookSet>(
                         let remaining_locals = local_func_inst.locals.iter().cloned();
                         let locals = Locals::new(params, remaining_locals);
 
-                        stack.push_stackframe(func_to_call_idx, func_to_call_ty, locals, wasm.pc);
+                        stack.push_stackframe(
+                            *current_module,
+                            func_to_call_idx,
+                            func_to_call_ty,
+                            locals,
+                            wasm.pc,
+                        );
+
+                        wasm.move_start_to(local_func_inst.code_expr)
+                            .unwrap_validated();
+                    }
+                    FuncInst::Imported(_imported_func_inst) => {
+                        let (next_module, next_func_idx) = lut
+                            .lookup(*current_module, func_to_call_idx)
+                            .expect("invalid state for lookup");
+
+                        let local_func_inst = modules[next_module].store.funcs[next_func_idx]
+                            .try_into_local()
+                            .unwrap();
+
+                        let remaining_locals = local_func_inst.locals.iter().cloned();
+                        let locals = Locals::new(params, remaining_locals);
+
+                        stack.push_stackframe(
+                            *current_module,
+                            func_to_call_idx,
+                            func_to_call_ty,
+                            locals,
+                            wasm.pc,
+                        );
+
+                        wasm = &mut modules[next_module].wasm_reader;
+                        *current_module = next_module;
 
                         wasm.move_start_to(local_func_inst.code_expr)
                             .unwrap_validated();
                     }
-                    FuncInst::Imported(imported_func_inst) => todo!(),
                 }
             }
             LOCAL_GET => {

src/execution/lut.rs

@@ -0,0 +1,77 @@
+use crate::{core::reader::types::export::ExportDesc, execution::execution_info::ExecutionInfo};
+use alloc::{collections::btree_map::BTreeMap, string::String, vec::Vec};
+
+pub struct LUT {
+    /// function_lut[local_module_idx][function_local_idx] = (foreign_module_idx, function_foreign_idx)
+    ///
+    /// - Module A imports a function "foo". Inside module A, the function has the index "function_local_idx". Module A
+    /// is assigned the index "local_module_idx".
+    /// - Module B exports a function "foo". Inside module B, the function has the index "function_foreign_idx". Module
+    /// B is assigned the index "foreign_module_idx".
+    function_lut: Vec<Vec<(usize, usize)>>,
+}
+
+impl LUT {
+    pub fn new(modules: &[ExecutionInfo], module_map: &BTreeMap<String, usize>) -> Option<Self> {
+        let mut function_lut = Vec::new();
+        for module in modules {
+            let module_lut = module
+                .store
+                .funcs
+                .iter()
+                .filter_map(|f| f.try_into_imported())
+                .map(|import| {
+                    Self::manual_lookup(
+                        modules,
+                        module_map,
+                        &import.module_name,
+                        &import.function_name,
+                    )
+                })
+                .collect::<Option<Vec<_>>>()?;
+
+            // TODO: what do we want to do if there is a missing import/export pair? Currently we fail the entire
+            // operation. Should it be a RuntimeError if said missing pair is called?
+
+            function_lut.push(module_lut);
+        }
+
+        Some(Self { function_lut })
+    }
+
+    pub fn lookup(&self, module_idx: usize, function_idx: usize) -> Option<(usize, usize)> {
+        self.function_lut
+            .get(module_idx)?
+            .get(function_idx)
+            .copied()
+    }
+
+    pub fn manual_lookup(
+        modules: &[ExecutionInfo],
+        module_map: &BTreeMap<String, usize>,
+        module_name: &str,
+        function_name: &str,
+    ) -> Option<(usize, usize)> {
+        let module_idx = module_map.get(module_name)?;
+        let module = &modules[*module_idx];
+
+        module
+            .store
+            .exports
+            .iter()
+            .filter_map(|export| {
+                if export.name == function_name {
+                    Some(&export.desc)
+                } else {
+                    None
+                }
+            })
+            .find_map(|desc| {
+                if let ExportDesc::FuncIdx(func_idx) = desc {
+                    Some((*module_idx, *func_idx))
+                } else {
+                    None
+                }
+            })
+    }
+}