return to compact side table, reduce runtime signature decoding, and misc optimizations

InvalidPathException · InvalidPathException · commit 4a0df270306e · 2025-09-23T02:22:34.000-04:00
diff --git a/src/instance.rs b/src/instance.rs
@@ -5,7 +5,7 @@ use crate::signature::{Signature, ValType};
 use crate::wasm_memory::WasmMemory;
 use crate::Module;
 use paste::paste;
-use std::cell::RefCell;
+use std::cell::{RefCell, Cell};
 use std::collections::HashMap;
 use std::rc::{Rc, Weak};
 
@@ -244,7 +244,7 @@ impl WasmTable {
 pub struct WasmGlobal {
     pub ty: ValType,
     pub mutable: bool,
-    pub value: WasmValue,
+    pub value: Cell<WasmValue>,
 }
 
 // --------------- Imports/Exports and Functions ---------------
@@ -419,7 +419,7 @@ impl Instance {
                     // evaluate constant initializer
                     let mut cpc = g.initializer_offset;
                     let val = Instance::eval_const(&module, &mut cpc, &inst.globals)?;
-                    inst.globals.push(Rc::new(RefCell::new(WasmGlobal { ty: g.ty, mutable: g.is_mutable, value: val })));
+                    inst.globals.push(Rc::new(RefCell::new(WasmGlobal { ty: g.ty, mutable: g.is_mutable, value: Cell::new(val) })));
                 }
             }
 
@@ -578,7 +578,7 @@ impl Instance {
                 0x42 => { let v: i64 = read_sleb128(bytes, pc)?; stack.push(WasmValue::from_i64(v)); }
                 0x43 => { let bits = u32::from_le_bytes(bytes[*pc..*pc+4].try_into().unwrap()); *pc += 4; stack.push(WasmValue::from_f32_bits(bits)); }
                 0x44 => { let bits = u64::from_le_bytes(bytes[*pc..*pc+8].try_into().unwrap()); *pc += 8; stack.push(WasmValue::from_f64_bits(bits)); }
-                0x23 => { let gi: u32 = read_leb128(bytes, pc)?; let g = gi as usize; if g >= globals.len() { return Err(Error::Validation(UNKNOWN_GLOBAL)); } stack.push(globals[g].borrow().value); }
+                0x23 => { let gi: u32 = read_leb128(bytes, pc)?; let g = gi as usize; if g >= globals.len() { return Err(Error::Validation(UNKNOWN_GLOBAL)); } stack.push(globals[g].borrow().value.get()); }
                 0x6a => { let b = stack.pop().unwrap().as_u32(); let a = stack.pop().unwrap().as_u32(); stack.push(WasmValue::from_u32(a.wrapping_add(b))); }
                 0x6b => { let b = stack.pop().unwrap().as_u32(); let a = stack.pop().unwrap().as_u32(); stack.push(WasmValue::from_u32(a.wrapping_sub(b))); }
                 0x6c => { let b = stack.pop().unwrap().as_u32(); let a = stack.pop().unwrap().as_u32(); stack.push(WasmValue::from_u32(a.wrapping_mul(b))); }
@@ -685,6 +685,7 @@ impl Instance {
         ctrl_bases: &mut Vec<usize>
     ) -> Result<(), Error> {
         let bytes = &self.module.bytes;
+        let mem_opt = self.memory.as_ref();
 
         macro_rules! next_op { () => {{ let byte = bytes[pc]; pc += 1; byte }} }
         macro_rules! pop_val { () => {{
@@ -891,7 +892,7 @@ impl Instance {
             let _align: u32 = read_leb128(bytes, &mut pc)?;
             let offset: u32 = read_leb128(bytes, &mut pc)?;
             let addr = pop_val!().as_u32();
-            let mem = self.memory.as_ref().ok_or_else(|| Error::Validation(UNKNOWN_MEMORY))?;
+            let mem = mem_opt.ok_or_else(|| Error::Validation(UNKNOWN_MEMORY))?;
             let v = mem.borrow().$method(addr, offset).map_err(|e| Error::Trap(e))?;
             let val = ($push)(v);
             stack.push(val);
@@ -902,7 +903,7 @@ impl Instance {
             let raw = pop_val!();
             let addr = pop_val!().as_u32();
             let val = ($from)(raw);
-            let mem = self.memory.as_ref().ok_or_else(|| Error::Validation(UNKNOWN_MEMORY))?;
+            let mem = mem_opt.ok_or_else(|| Error::Validation(UNKNOWN_MEMORY))?;
             mem.borrow_mut().$method(addr, offset, val).map_err(|e| Error::Trap(e))?;
         }}}
 
@@ -912,36 +913,36 @@ impl Instance {
                 0x00 => return Err(Error::Trap(UNREACHABLE)),
                 0x01 | 0xbc | 0xbd | 0xbe | 0xbf => {} // nop and reinterprets (no-op on raw bits)
                 0x02 => { // block
-                    let sig = Signature::read(&self.module.types, bytes, &mut pc)?;
-                    let block_end = *self.module.block_ends.get(&pc).unwrap();
+                    let entry = *self.module.side_table.get(&(pc as u32)).unwrap();
+                    pc = entry.body_pc as usize;
                     control.push(ControlFrame {
-                        stack_len: stack.len() - sig.params.len(),
-                        dest_pc: block_end,
-                        arity: sig.result.is_some() as u32,
-                        has_result: sig.result.is_some()
+                        stack_len: stack.len() - (entry.params_len as usize),
+                        dest_pc: entry.end_pc as usize,
+                        arity: entry.has_result as u32,
+                        has_result: entry.has_result
                     });
                 }
                 0x03 => { // loop
                     let loop_op_pc = pc - 1;
-                    let sig = Signature::read(&self.module.types, bytes, &mut pc)?;
+                    let entry = *self.module.side_table.get(&(pc as u32)).unwrap();
+                    pc = entry.body_pc as usize;
                     control.push(ControlFrame {
-                        stack_len: stack.len() - sig.params.len(),
+                        stack_len: stack.len() - (entry.params_len as usize),
                         dest_pc: loop_op_pc,
-                        arity: sig.params.len() as u32,
-                        has_result: sig.result.is_some()
+                        arity: entry.params_len as u32,
+                        has_result: entry.has_result
                     });
                 }
                 0x04 => { // if
-                    let sig = Signature::read(&self.module.types, bytes, &mut pc)?;
+                    let entry = *self.module.side_table.get(&(pc as u32)).unwrap();
                     let cond = pop_val!().as_u32();
-                    let if_jump = self.module.if_jumps.get(&pc).unwrap();
                     control.push(ControlFrame {
-                        stack_len: stack.len() - sig.params.len(),
-                        dest_pc: if_jump.end_offset,
-                        arity: sig.result.is_some() as u32,
-                        has_result: sig.result.is_some()
+                        stack_len: stack.len() - (entry.params_len as usize),
+                        dest_pc: entry.end_pc as usize,
+                        arity: entry.has_result as u32,
+                        has_result: entry.has_result
                     });
-                    if cond == 0 { pc = if_jump.else_offset; }
+                    pc = if cond == 0 { entry.else_pc as usize } else { entry.body_pc as usize };
                 }
                 0x05 => { // else
                     let _ = Instance::branch(&mut pc, stack, control, 0);
@@ -1009,6 +1010,9 @@ impl Instance {
                 }
                 // Call instructions
                 0x10 => { // call
+                    // direct calls are fully type-checked at validation time; no
+                    // structural type check is required here. we only use the
+                    // runtime signature for fast param/result counts to set up frames.
                     let fi: u32 = read_leb128(bytes, &mut pc)?;
                     let f = &self.functions[fi as usize];
                     
@@ -1047,6 +1051,8 @@ impl Instance {
                     }
                 }
                 0x11 => { // call_indirect
+                    // Indirect calls must enforce params at runtime
+                    // Here we must parse the indices
                     let type_idx: u32 = read_leb128(bytes, &mut pc)?;
                     pc += 1; // Skip the zero flag
                     let elem_idx = match stack.pop() {
@@ -1206,7 +1212,7 @@ impl Instance {
                     if gi as usize >= self.globals.len() {
                         return Err(Error::Trap(UNKNOWN_GLOBAL));
                     }
-                    stack.push(self.globals[gi as usize].borrow().value);
+                    stack.push(self.globals[gi as usize].borrow().value.get());
                 }
                 0x24 => { // global.set
                     let gi: u32 = read_leb128(bytes, &mut pc)?;
@@ -1217,7 +1223,7 @@ impl Instance {
                         Some(v) => v,
                         None => return Err(Error::Trap(STACK_UNDERFLOW))
                     };
-                    self.globals[gi as usize].borrow_mut().value = val;
+                    self.globals[gi as usize].borrow().value.set(val);
                 }
                 // Memory instructions - loads
                 0x28 => { load!(load_u32, |v: u32| WasmValue::from_u32(v)); }
diff --git a/src/module.rs b/src/module.rs
@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use nohash_hasher::BuildNoHashHasher;
 use std::rc::Rc;
 
 use crate::byte_iter::*;
@@ -67,8 +68,14 @@ pub struct Export { pub extern_type: ExternType, pub idx: u32 }
 #[derive(Clone)]
 pub struct DataSegment { pub data_range: std::ops::Range<usize>, pub initializer_offset: usize }
 
-#[derive(Clone, Copy)]
-pub struct IfJump { pub else_offset: usize, pub end_offset: usize }
+#[derive(Clone, Copy, Default)]
+pub struct SideTableEntry {
+    pub params_len: u16,
+    pub has_result: bool,
+    pub body_pc: u32,   // start pc of block body (after the sig)
+    pub end_pc: u32,    // end pc of the construct
+    pub else_pc: u32,   // else start pc or end_pc if no else
+}
 
 // ---------------- Module Structure ----------------
 #[derive(Default)]
@@ -86,8 +93,7 @@ pub struct Module {
     pub functions: Vec<Function>,
     pub n_data: u32,
     pub data_segments: Vec<DataSegment>,
-    pub if_jumps: HashMap<usize, IfJump>,
-    pub block_ends: HashMap<usize, usize>,
+    pub side_table: HashMap<u32, SideTableEntry, BuildNoHashHasher<u32>>,
 }
 
 impl Module {
diff --git a/src/validator.rs b/src/validator.rs
@@ -24,6 +24,7 @@ pub struct ControlFrame {
     pub height: usize,
     pub unreachable: bool,
     pub control_type: ControlType,
+    pub sig_pc: usize,
 }
 
 // ---------------- ValidatorStack for Type Checking ----------------
@@ -78,12 +79,13 @@ impl ValidatorStack {
         Ok(popped)
     }
 
-    pub fn push_ctrl(&mut self, sig: Signature, control_type: ControlType) -> Result<(), Error> {
+    pub fn push_ctrl(&mut self, sig: Signature, control_type: ControlType, sig_pc: usize) -> Result<(), Error> {
         let frame = ControlFrame {
             sig: sig.clone(),
             height: self.val_stack.len(),
             unreachable: false,
             control_type,
+            sig_pc,
         };
         self.ctrl_stack.push(frame);
         self.push_vals(&sig.params);
@@ -188,7 +190,8 @@ impl<'a> Validator<'a> {
             sig: func.ty.clone(), 
             height: func.ty.params.len(),  // Stack height after params
             unreachable: false,
-            control_type: ControlType::Function 
+            control_type: ControlType::Function,
+            sig_pc: func.body.start.saturating_sub(1),
         });
 
         // Validation loop
@@ -229,27 +232,58 @@ fn validate_nop(_: &mut Module, _: &mut ByteIter, _: &Function, _: &mut Validato
     Ok(Action::Continue)
 }
 
-fn validate_block(m: &mut Module, it: &mut ByteIter, _: &Function, vs: &mut ValidatorStack) -> Result<Action, Error> {
+    fn validate_block(m: &mut Module, it: &mut ByteIter, _: &Function, vs: &mut ValidatorStack) -> Result<Action, Error> {
+    let sig_pc = it.cur();
     let sig = Signature::read(&m.types, &m.bytes, &mut it.idx)?;
     let block_start = it.cur();
     vs.pop_vals(&sig.params)?;
-    vs.push_ctrl(sig, ControlType::Block { start: block_start })?;
+    let params_len = sig.params.len() as u16;
+    let has_result = sig.result.is_some();
+    vs.push_ctrl(sig, ControlType::Block { start: block_start }, sig_pc)?;
+    m.side_table.entry(sig_pc as u32).or_insert(SideTableEntry {
+        params_len,
+        has_result,
+        body_pc: block_start as u32,
+        end_pc: 0,
+        else_pc: 0,
+    });
     Ok(Action::Continue)
 }
 
-fn validate_loop(m: &mut Module, it: &mut ByteIter, _: &Function, vs: &mut ValidatorStack) -> Result<Action, Error> {
+    fn validate_loop(m: &mut Module, it: &mut ByteIter, _: &Function, vs: &mut ValidatorStack) -> Result<Action, Error> {
+    let sig_pc = it.cur();
     let sig = Signature::read(&m.types, &m.bytes, &mut it.idx)?;
+    let loop_body_pc = it.cur(); // body starts here
     vs.pop_vals(&sig.params)?;
-    vs.push_ctrl(sig, ControlType::Loop)?;
+    let params_len = sig.params.len() as u16;
+    let has_result = sig.result.is_some();
+    vs.push_ctrl(sig, ControlType::Loop, sig_pc)?;
+    m.side_table.entry(sig_pc as u32).or_insert(SideTableEntry {
+        params_len,
+        has_result,
+        body_pc: loop_body_pc as u32,
+        end_pc: 0,
+        else_pc: loop_body_pc as u32, // for loop back-edge
+    });
     Ok(Action::Continue)
 }
 
-fn validate_if(m: &mut Module, it: &mut ByteIter, _: &Function, vs: &mut ValidatorStack) -> Result<Action, Error> {
+    fn validate_if(m: &mut Module, it: &mut ByteIter, _: &Function, vs: &mut ValidatorStack) -> Result<Action, Error> {
+    let sig_pc = it.cur();
     let sig = Signature::read(&m.types, &m.bytes, &mut it.idx)?;
     vs.pop_val_expect(ValType::I32)?;
     vs.pop_vals(&sig.params)?;
-    let if_start = it.cur();
-    vs.push_ctrl(sig, ControlType::If { start: if_start })?;
+    let if_body_pc = it.cur();
+    let params_len = sig.params.len() as u16;
+    let has_result = sig.result.is_some();
+    vs.push_ctrl(sig, ControlType::If { start: if_body_pc }, sig_pc)?;
+    m.side_table.entry(sig_pc as u32).or_insert(SideTableEntry {
+        params_len,
+        has_result,
+        body_pc: if_body_pc as u32,
+        end_pc: 0,
+        else_pc: 0, // filled at else/end
+    });
     Ok(Action::Continue)
 }
 
@@ -288,6 +322,7 @@ fn validate_else(_: &mut Module, it: &mut ByteIter, _: &Function, vs: &mut Valid
         height: frame.height,
         unreachable: false,
         control_type: new_control_type,
+        sig_pc: frame.sig_pc,
     });
     vs.push_vals(&params);
     Ok(Action::Continue)
@@ -317,22 +352,34 @@ fn validate_end(m: &mut Module, it: &mut ByteIter, f: &Function, vs: &mut Valida
     
     // Handle jump offset tracking
     match frame.control_type {
-        ControlType::Block { start } => {
-            m.block_ends.insert(start, it.cur());
+        ControlType::Block { .. } => {
+            let k = frame.sig_pc as u32;
+            let mut updated = m.side_table.get(&k).copied().unwrap_or_default();
+            updated.end_pc = it.cur() as u32;
+            updated.else_pc = updated.end_pc;
+            m.side_table.insert(k, updated);
         }
         ControlType::Loop => {}
-        ControlType::If { start } => {
+        ControlType::If { .. } => {
             // For if without else, params must equal results
             let results_as_vec: Vec<ValType> = frame.sig.result.into_iter().collect();
             if frame.sig.params != results_as_vec {
                 return Err(Validation(TYPE_MISMATCH));
             }
             let else_off = it.cur() - 1;
             let end_off = it.cur();
-            m.if_jumps.insert(start, IfJump { else_offset: else_off, end_offset: end_off });
+            let k = frame.sig_pc as u32;
+            let mut updated = m.side_table.get(&k).copied().unwrap_or_default();
+            updated.end_pc = end_off as u32;
+            updated.else_pc = else_off as u32;
+            m.side_table.insert(k, updated);
         }
-        ControlType::IfElse { if_start, else_start } => {
-            m.if_jumps.insert(if_start, IfJump { else_offset: else_start, end_offset: it.cur() });
+        ControlType::IfElse { else_start, .. } => {
+            let k = frame.sig_pc as u32;
+            let mut updated = m.side_table.get(&k).copied().unwrap_or_default();
+            updated.end_pc = it.cur() as u32;
+            updated.else_pc = else_start as u32;
+            m.side_table.insert(k, updated);
         }
         ControlType::Function => {}
     }
diff --git a/tests/spec_tests.rs b/tests/spec_tests.rs
@@ -1,4 +1,4 @@
-use std::{collections::HashMap, fs, path::Path, process::Command, rc::{Rc}, cell::RefCell, env};
+use std::{collections::HashMap, fs, path::Path, process::Command, rc::{Rc}, cell::{RefCell, Cell}, env};
 use serde::Deserialize;
 use wagmi::{Module, Instance, Imports, ExportValue, WasmValue, WasmGlobal, WasmTable, WasmMemory, RuntimeFunction, RuntimeSignature, ValType, Error, Signature};
 
@@ -107,16 +107,16 @@ fn spectest_exports() -> HashMap<String, ExportValue> {
     let mut exports = HashMap::new();
 
     exports.insert("global_i32".into(), ExportValue::Global(Rc::new(RefCell::new(
-        WasmGlobal { ty: ValType::I32, mutable: false, value: WasmValue::from_u32(666) }
+        WasmGlobal { ty: ValType::I32, mutable: false, value: Cell::new(WasmValue::from_u32(666)) }
     ))));
     exports.insert("global_i64".into(), ExportValue::Global(Rc::new(RefCell::new(
-        WasmGlobal { ty: ValType::I64, mutable: false, value: WasmValue::from_u64(666) }
+        WasmGlobal { ty: ValType::I64, mutable: false, value: Cell::new(WasmValue::from_u64(666)) }
     ))));
     exports.insert("global_f32".into(), ExportValue::Global(Rc::new(RefCell::new(
-        WasmGlobal { ty: ValType::F32, mutable: false, value: WasmValue::from_f32(666.6) }
+        WasmGlobal { ty: ValType::F32, mutable: false, value: Cell::new(WasmValue::from_f32(666.6)) }
     ))));
     exports.insert("global_f64".into(), ExportValue::Global(Rc::new(RefCell::new(
-        WasmGlobal { ty: ValType::F64, mutable: false, value: WasmValue::from_f64(666.6) }
+        WasmGlobal { ty: ValType::F64, mutable: false, value: Cell::new(WasmValue::from_f64(666.6)) }
     ))));
 
     exports.insert("table".into(), ExportValue::Table(Rc::new(RefCell::new(
@@ -157,7 +157,7 @@ fn exec_action(instances: &HashMap<String, Rc<Instance>>, action: &Act) -> Resul
     
     match action {
         Act::Get { .. } => match export {
-            ExportValue::Global(g) => Ok(vec![g.borrow().value]),
+            ExportValue::Global(g) => Ok(vec![g.borrow().value.get()]),
             _ => Err(Error::Trap("not a global"))
         },
         Act::Invoke { .. } => match export {
