fix: per module sidetables

Signed-off-by: Cem Onem <[email protected]>

src/execution/interpreter_loop.rs

@@ -54,10 +54,8 @@ pub(super) fn run<H: HookSet>(
     // Start reading the function's instructions
     let mut wasm = &mut modules[*current_module_idx].wasm_reader;
 
-    // the sidetable and stp for this function, stp will reset to 0 every call
-    // since function instances have their own sidetable.
-    let mut current_sidetable: &Sidetable = &func_inst.sidetable;
-    let mut stp = 0;
+    let mut current_sidetable = &modules[*current_module_idx].store.sidetable;
+    let mut stp = func_inst.stp;
 
     // unwrap is sound, because the validation assures that the function points to valid subslice of the WASM binary
     wasm.move_start_to(func_inst.code_expr).unwrap();
@@ -109,14 +107,7 @@ pub(super) fn run<H: HookSet>(
                 wasm.pc = maybe_return_address;
                 stp = maybe_return_stp;
 
-                current_sidetable = &modules[return_module]
-                    .store
-                    .funcs
-                    .get(stack.current_stackframe().func_idx)
-                    .unwrap_validated()
-                    .try_into_local()
-                    .unwrap_validated()
-                    .sidetable;
+                current_sidetable = &modules[return_module].store.sidetable;
 
                 *current_module_idx = return_module;
             }
@@ -209,8 +200,7 @@ pub(super) fn run<H: HookSet>(
                         wasm.move_start_to(local_func_inst.code_expr)
                             .unwrap_validated();
 
-                        stp = 0;
-                        current_sidetable = &local_func_inst.sidetable;
+                        stp = local_func_inst.stp;
                     }
                     FuncInst::Imported(_imported_func_inst) => {
                         let (next_module, next_func_idx) = lut
@@ -239,8 +229,8 @@ pub(super) fn run<H: HookSet>(
                         wasm.move_start_to(local_func_inst.code_expr)
                             .unwrap_validated();
 
-                        stp = 0;
-                        current_sidetable = &local_func_inst.sidetable;
+                        stp = local_func_inst.stp;
+                        current_sidetable = &modules[next_module].store.sidetable;
                     }
                 }
             }
@@ -309,8 +299,7 @@ pub(super) fn run<H: HookSet>(
                         wasm.move_start_to(local_func_inst.code_expr)
                             .unwrap_validated();
 
-                        stp = 0;
-                        current_sidetable = &local_func_inst.sidetable;
+                        stp = local_func_inst.stp;
                     }
                     FuncInst::Imported(_imported_func_inst) => {
                         let (next_module, next_func_idx) = lut
@@ -341,8 +330,8 @@ pub(super) fn run<H: HookSet>(
                         wasm.move_start_to(local_func_inst.code_expr)
                             .unwrap_validated();
 
-                        stp = 0;
-                        current_sidetable = &local_func_inst.sidetable;
+                        stp = local_func_inst.stp;
+                        current_sidetable = &modules[next_module].store.sidetable;
                     }
                 }
             }

src/execution/mod.rs

@@ -422,9 +422,9 @@ where
             let mut wasm_reader = WasmReader::new(validation_info.wasm);
 
             let functions = validation_info.functions.iter();
-            let func_blocks = validation_info.func_blocks.iter();
+            let func_blocks_stps = validation_info.func_blocks_stps.iter();
 
-            let local_function_inst = functions.zip(func_blocks).map(|(ty, (func, sidetable))| {
+            let local_function_inst = functions.zip(func_blocks_stps).map(|(ty, (func, stp))| {
                 wasm_reader
                     .move_start_to(*func)
                     .expect("function index to be in the bounds of the WASM binary");
@@ -441,8 +441,7 @@ where
                     ty: *ty,
                     locals,
                     code_expr,
-                    // TODO fix this ugly clone
-                    sidetable: sidetable.clone(),
+                    stp: *stp,
                 })
             });
 
@@ -643,6 +642,8 @@ where
             elements,
             passive_elem_indexes,
             exports,
+            // TODO is this ok? possible multiple instantiations of the store from same code forces us a clone
+            sidetable: validation_info.sidetable.clone(),
         })
     }
 }

src/execution/store.rs

@@ -26,6 +26,7 @@ pub struct Store {
     pub elements: Vec<ElemInst>,
     pub passive_elem_indexes: Vec<usize>,
     pub exports: Vec<Export>,
+    pub sidetable: Sidetable,
 }
 
 #[derive(Debug)]
@@ -39,7 +40,8 @@ pub struct LocalFuncInst {
     pub ty: TypeIdx,
     pub locals: Vec<ValType>,
     pub code_expr: Span,
-    pub sidetable: Sidetable,
+    /// where stp should point to when pc points to the beginning of the function
+    pub stp: usize,
 }
 
 #[derive(Debug)]

src/validation/code.rs

@@ -30,11 +30,11 @@ pub fn validate_code_section(
     tables: &[TableType],
     elements: &[ElemType],
     referenced_functions: &BTreeSet<u32>,
-) -> Result<Vec<(Span, Sidetable)>> {
+    sidetable: &mut Sidetable,
+) -> Result<Vec<(Span, usize)>> {
     assert_eq!(section_header.ty, SectionTy::Code);
 
-    // TODO replace with single sidetable per module
-    let code_block_spans_sidetables = wasm.read_vec_enumerated(|wasm, idx| {
+    let code_block_spans_stps = wasm.read_vec_enumerated(|wasm, idx| {
         // We need to offset the index by the number of functions that were
         // imported. Imported functions always live at the start of the index
         // space.
@@ -51,13 +51,13 @@ pub fn validate_code_section(
             params.chain(declared_locals).collect::<Vec<ValType>>()
         };
 
-        let mut stack = ValidationStack::new_for_func(func_ty);
-        let mut sidetable: Sidetable = Sidetable::default();
+        let stp = sidetable.len();
 
+        let mut stack = ValidationStack::new_for_func(func_ty);
         read_instructions(
             wasm,
             &mut stack,
-            &mut sidetable,
+            sidetable,
             &locals,
             globals,
             fn_types,
@@ -76,15 +76,15 @@ pub fn validate_code_section(
             )
         }
 
-        Ok((func_block, sidetable))
+        Ok((func_block, stp))
     })?;
 
     trace!(
         "Read code section. Found {} code blocks",
-        code_block_spans_sidetables.len()
+        code_block_spans_stps.len()
     );
 
-    Ok(code_block_spans_sidetables)
+    Ok(code_block_spans_stps)
 }
 
 pub fn read_declared_locals(wasm: &mut WasmReader) -> Result<Vec<ValType>> {

src/validation/mod.rs

@@ -33,11 +33,12 @@ pub struct ValidationInfo<'bytecode> {
     #[allow(dead_code)]
     pub(crate) exports: Vec<Export>,
     /// Each block contains the validated code section and the generated sidetable
-    pub(crate) func_blocks: Vec<(Span, Sidetable)>,
+    pub(crate) func_blocks_stps: Vec<(Span, usize)>,
     pub(crate) data: Vec<DataSegment>,
     /// The start function which is automatically executed during instantiation
     pub(crate) start: Option<FuncIdx>,
     pub(crate) elements: Vec<ElemType>,
+    pub(crate) sidetable: Sidetable,
 }
 
 pub fn validate(wasm: &[u8]) -> Result<ValidationInfo> {
@@ -169,26 +170,27 @@ pub fn validate(wasm: &[u8]) -> Result<ValidationInfo> {
 
     while (skip_section(&mut wasm, &mut header)?).is_some() {}
 
-    let func_blocks_sidetables =
-        handle_section(&mut wasm, &mut header, SectionTy::Code, |wasm, h| {
-            code::validate_code_section(
-                wasm,
-                h,
-                &types,
-                &all_functions,
-                imported_functions.count(),
-                &globals,
-                &memories,
-                &data_count,
-                &tables,
-                &elements,
-                &referenced_functions,
-            )
-        })?
-        .unwrap_or_default();
+    let mut sidetable = Sidetable::new();
+    let func_blocks_stps = handle_section(&mut wasm, &mut header, SectionTy::Code, |wasm, h| {
+        code::validate_code_section(
+            wasm,
+            h,
+            &types,
+            &all_functions,
+            imported_functions.count(),
+            &globals,
+            &memories,
+            &data_count,
+            &tables,
+            &elements,
+            &referenced_functions,
+            &mut sidetable,
+        )
+    })?
+    .unwrap_or_default();
 
     assert_eq!(
-        func_blocks_sidetables.len(),
+        func_blocks_stps.len(),
         local_functions.len(),
         "these should be equal"
     ); // TODO check if this is in the spec
@@ -222,10 +224,11 @@ pub fn validate(wasm: &[u8]) -> Result<ValidationInfo> {
         memories,
         globals,
         exports,
-        func_blocks: func_blocks_sidetables,
+        func_blocks_stps,
         data: data_section,
         start,
         elements,
+        sidetable,
     })
 }