chore: fix stable build, move BranchTableTarget to FunctionData

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

6 files changed, 114 insertions, 100 deletions

diff --git a/crates/parser/src/module.rs b/crates/parser/src/module.rs
index fdc02ab..29e556c 100644
--- a/crates/parser/src/module.rs
+++ b/crates/parser/src/module.rs
@@ -196,7 +196,7 @@ impl ModuleReader {
             .into_iter()
             .zip(self.code_type_addrs)
             .enumerate()
-            .map(|(func_idx, ((instructions, data, locals), ty_idx))| {
+            .map(|(func_idx, ((instructions, mut data, locals), ty_idx))| {
                 let ty = self.func_types.get(ty_idx as usize).expect("No func type for func, this is a bug").clone();
                 let params = ValueCountsSmall::from(&ty.params);
                 let locals = ValueCountsSmall {
@@ -206,7 +206,7 @@ impl ModuleReader {
                     cref: u16::try_from(locals.cref).unwrap_or_else(|_| unreachable!("local count exceeds u16")),
                 };
                 let self_func_addr = imported_func_count + func_idx as u32;
-                let instructions = optimize::optimize_instructions(instructions, self_func_addr, options);
+                let instructions = optimize::optimize_instructions(instructions, &mut data, self_func_addr, options);
 
                 WasmFunction { instructions: ArcSlice::from(instructions), data, locals, params, ty }
             })
diff --git a/crates/parser/src/optimize.rs b/crates/parser/src/optimize.rs
index 887ea04..1369d7a 100644
--- a/crates/parser/src/optimize.rs
+++ b/crates/parser/src/optimize.rs
@@ -1,15 +1,16 @@
 use crate::ParserOptions;
 use alloc::vec::Vec;
-use tinywasm_types::Instruction;
+use tinywasm_types::{Instruction, WasmFunctionData};
 
 pub(crate) fn optimize_instructions(
     mut instructions: Vec<Instruction>,
+    function_data: &mut WasmFunctionData,
     self_func_addr: u32,
     options: &ParserOptions,
 ) -> Vec<Instruction> {
     rewrite(&mut instructions, self_func_addr);
     if options.dce {
-        dce(&mut instructions);
+        dce(&mut instructions, function_data);
     }
     instructions
 }
@@ -135,29 +136,32 @@ fn rewrite(instructions: &mut [Instruction], self_func_addr: u32) {
                     instructions[read] = Instruction::Nop;
                 }
             }
-            Instruction::LocalGet64(dst)
+            Instruction::LocalGet64(dst) => {
                 if read > 0
                     && let Instruction::LocalSet64(src) = instructions[read - 1]
-                    && src == dst =>
-            {
-                instructions[read - 1] = Instruction::LocalTee64(src);
-                instructions[read] = Instruction::Nop;
+                    && src == dst
+                {
+                    instructions[read - 1] = Instruction::LocalTee64(src);
+                    instructions[read] = Instruction::Nop;
+                }
             }
-            Instruction::LocalGet128(dst)
+            Instruction::LocalGet128(dst) => {
                 if read > 0
                     && let Instruction::LocalSet128(src) = instructions[read - 1]
-                    && src == dst =>
-            {
-                instructions[read - 1] = Instruction::LocalTee128(src);
-                instructions[read] = Instruction::Nop;
+                    && src == dst
+                {
+                    instructions[read - 1] = Instruction::LocalTee128(src);
+                    instructions[read] = Instruction::Nop;
+                }
             }
-            Instruction::LocalGetRef(dst)
+            Instruction::LocalGetRef(dst) => {
                 if read > 0
                     && let Instruction::LocalSetRef(src) = instructions[read - 1]
-                    && src == dst =>
-            {
-                instructions[read - 1] = Instruction::LocalTeeRef(src);
-                instructions[read] = Instruction::Nop;
+                    && src == dst
+                {
+                    instructions[read - 1] = Instruction::LocalTeeRef(src);
+                    instructions[read] = Instruction::Nop;
+                }
             }
 
             Instruction::LocalSet32(dst) => {
@@ -230,19 +234,23 @@ fn rewrite(instructions: &mut [Instruction], self_func_addr: u32) {
                     instructions[read] = Instruction::LocalAddConst64(dst, c);
                 }
             }
-            Instruction::LocalSet128(dst)
+            Instruction::LocalSet128(dst) => {
                 if read > 0
-                    && let Instruction::LocalGet128(src) = instructions[read - 1] =>
-            {
-                instructions[read - 1] = Instruction::Nop;
-                instructions[read] = if src == dst { Instruction::Nop } else { Instruction::LocalCopy128(src, dst) };
+                    && let Instruction::LocalGet128(src) = instructions[read - 1]
+                {
+                    instructions[read - 1] = Instruction::Nop;
+                    instructions[read] =
+                        if src == dst { Instruction::Nop } else { Instruction::LocalCopy128(src, dst) };
+                }
             }
-            Instruction::LocalSetRef(dst)
+            Instruction::LocalSetRef(dst) => {
                 if read > 0
-                    && let Instruction::LocalGetRef(src) = instructions[read - 1] =>
-            {
-                instructions[read - 1] = Instruction::Nop;
-                instructions[read] = if src == dst { Instruction::Nop } else { Instruction::LocalCopyRef(src, dst) };
+                    && let Instruction::LocalGetRef(src) = instructions[read - 1]
+                {
+                    instructions[read - 1] = Instruction::Nop;
+                    instructions[read] =
+                        if src == dst { Instruction::Nop } else { Instruction::LocalCopyRef(src, dst) };
+                }
             }
 
             Instruction::LocalTee32(dst) => {
@@ -273,55 +281,61 @@ fn rewrite(instructions: &mut [Instruction], self_func_addr: u32) {
                 }
                 _ => {}
             },
-            Instruction::LocalTee128(dst)
+            Instruction::LocalTee128(dst) => {
                 if read > 0
                     && let Instruction::LocalGet128(src) = instructions[read - 1]
-                    && src == dst =>
-            {
-                instructions[read] = Instruction::Nop;
+                    && src == dst
+                {
+                    instructions[read] = Instruction::Nop;
+                }
             }
-            Instruction::LocalTeeRef(dst)
+            Instruction::LocalTeeRef(dst) => {
                 if read > 0
                     && let Instruction::LocalGetRef(src) = instructions[read - 1]
-                    && src == dst =>
-            {
-                instructions[read] = Instruction::Nop;
+                    && src == dst
+                {
+                    instructions[read] = Instruction::Nop;
+                }
             }
 
-            Instruction::Drop32
+            Instruction::Drop32 => {
                 if read > 0
-                    && let Instruction::LocalTee32(local) = instructions[read - 1] =>
-            {
-                instructions[read - 1] = Instruction::LocalSet32(local);
-                instructions[read] = Instruction::Nop;
+                    && let Instruction::LocalTee32(local) = instructions[read - 1]
+                {
+                    instructions[read - 1] = Instruction::LocalSet32(local);
+                    instructions[read] = Instruction::Nop;
+                }
             }
-            Instruction::Drop64
+            Instruction::Drop64 => {
                 if read > 0
-                    && let Instruction::LocalTee64(local) = instructions[read - 1] =>
-            {
-                instructions[read - 1] = Instruction::LocalSet64(local);
-                instructions[read] = Instruction::Nop;
+                    && let Instruction::LocalTee64(local) = instructions[read - 1]
+                {
+                    instructions[read - 1] = Instruction::LocalSet64(local);
+                    instructions[read] = Instruction::Nop;
+                }
             }
-            Instruction::Drop128
+            Instruction::Drop128 => {
                 if read > 0
-                    && let Instruction::LocalTee128(local) = instructions[read - 1] =>
-            {
-                instructions[read - 1] = Instruction::LocalSet128(local);
-                instructions[read] = Instruction::Nop;
+                    && let Instruction::LocalTee128(local) = instructions[read - 1]
+                {
+                    instructions[read - 1] = Instruction::LocalSet128(local);
+                    instructions[read] = Instruction::Nop;
+                }
             }
-            Instruction::DropRef
+            Instruction::DropRef => {
                 if read > 0
-                    && let Instruction::LocalTeeRef(local) = instructions[read - 1] =>
-            {
-                instructions[read - 1] = Instruction::LocalSetRef(local);
-                instructions[read] = Instruction::Nop;
+                    && let Instruction::LocalTeeRef(local) = instructions[read - 1]
+                {
+                    instructions[read - 1] = Instruction::LocalSetRef(local);
+                    instructions[read] = Instruction::Nop;
+                }
             }
             _ => {}
         }
     }
 }
 
-fn dce(instructions: &mut Vec<Instruction>) {
+fn dce(instructions: &mut Vec<Instruction>, function_data: &mut WasmFunctionData) {
     let old_len = instructions.len();
     if old_len == 0 {
         return;
@@ -340,13 +354,21 @@ fn dce(instructions: &mut Vec<Instruction>) {
     }
 
     let compacted_len = old_len as u32 - removed_total;
+
+    function_data.branch_table_targets.iter_mut().for_each(|ip| {
+        let old_target = *ip as usize;
+        if old_target <= old_len {
+            *ip -= removed_before[old_target];
+            debug_assert!(*ip < compacted_len, "remapped jump target points past end of function");
+        }
+    });
+
     instructions.retain_mut(|instr| {
         let ip = match instr {
             Instruction::Jump(ip)
             | Instruction::JumpIfZero(ip)
             | Instruction::JumpIfNonZero(ip)
-            | Instruction::BranchTableTarget(ip)
-            | Instruction::BranchTable(ip, _) => ip,
+            | Instruction::BranchTable(ip, _, _) => ip,
             _ => return !matches!(instr, Instruction::Nop),
         };
 
diff --git a/crates/parser/src/visit.rs b/crates/parser/src/visit.rs
index 09e3fae..076f7c1 100644
--- a/crates/parser/src/visit.rs
+++ b/crates/parser/src/visit.rs
@@ -32,6 +32,21 @@ struct LoweringCtx {
     branch_jumps: Vec<usize>,
 }
 
+#[derive(Default)]
+struct FunctionDataBuilder {
+    v128_constants: Vec<i128>,
+    branch_table_targets: Vec<u32>,
+}
+
+impl FunctionDataBuilder {
+    fn finish(self) -> WasmFunctionData {
+        WasmFunctionData {
+            v128_constants: self.v128_constants.into_boxed_slice(),
+            branch_table_targets: self.branch_table_targets.into_boxed_slice(),
+        }
+    }
+}
+
 struct ValidateThenVisit<'a, R: WasmModuleResources>(usize, &'a mut FunctionBuilder<R>);
 
 macro_rules! validate_then_visit {
@@ -75,11 +90,7 @@ pub(crate) fn process_operators_and_validate<R: WasmModuleResources>(
         return Err(builder.errors.remove(0));
     }
 
-    Ok((
-        builder.instructions,
-        WasmFunctionData { v128_constants: builder.v128_constants.into_boxed_slice() },
-        builder.validator.into_allocations(),
-    ))
+    Ok((builder.instructions, builder.data.finish(), builder.validator.into_allocations()))
 }
 
 macro_rules! define_operand {
@@ -135,7 +146,7 @@ macro_rules! define_mem_operands_simd_lane {
 pub(crate) struct FunctionBuilder<R: WasmModuleResources> {
     validator: FuncValidator<R>,
     instructions: Vec<Instruction>,
-    v128_constants: Vec<i128>,
+    data: FunctionDataBuilder,
     ctx_stack: Vec<LoweringCtx>,
     local_addr_map: Vec<u32>,
     errors: Vec<crate::ParseError>,
@@ -390,14 +401,8 @@ impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuild
         let target_depths: Vec<u32> = ts;
 
         let header_ip = self.instructions.len();
-        self.instructions.push(Instruction::BranchTable(0, len));
-
-        let target_table_ip = self.instructions.len();
-        for _ in 0..len {
-            self.instructions.push(Instruction::BranchTableTarget(0));
-        }
-        let default_target_ip = self.instructions.len();
-        self.instructions.push(Instruction::BranchTableTarget(0));
+        let branch_table_start = self.data.branch_table_targets.len() as u32;
+        self.instructions.push(Instruction::BranchTable(0, branch_table_start, len));
 
         let mut seen = alloc::collections::BTreeMap::<u32, usize>::new();
         struct PadInfo {
@@ -418,18 +423,13 @@ impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuild
             pads.push(PadInfo { depth, pad_start, jump_or_ret_ip, is_return });
         }
 
-        for (i, &depth) in target_depths.iter().enumerate() {
+        for &depth in &target_depths {
             let pad_idx = seen[&depth];
-            if let Instruction::BranchTableTarget(ip) = &mut self.instructions[target_table_ip + i] {
-                *ip = pads[pad_idx].pad_start as u32;
-            }
+            self.data.branch_table_targets.push(pads[pad_idx].pad_start as u32);
         }
 
         let default_pad_idx = seen[&default_depth];
-        if let Instruction::BranchTableTarget(ip) = &mut self.instructions[default_target_ip] {
-            *ip = pads[default_pad_idx].pad_start as u32;
-        }
-        if let Instruction::BranchTable(default_ip, _) = &mut self.instructions[header_ip] {
+        if let Instruction::BranchTable(default_ip, _, _) = &mut self.instructions[header_ip] {
             *default_ip = pads[default_pad_idx].pad_start as u32;
         }
 
@@ -570,13 +570,13 @@ impl<R: WasmModuleResources> wasmparser::VisitSimdOperator<'_> for FunctionBuild
     }
 
     fn visit_i8x16_shuffle(&mut self, lanes: [u8; 16]) -> Self::Output {
-        self.instructions.push(Instruction::I8x16Shuffle(self.v128_constants.len() as u32));
-        self.v128_constants.push(i128::from_le_bytes(lanes));
+        self.instructions.push(Instruction::I8x16Shuffle(self.data.v128_constants.len() as u32));
+        self.data.v128_constants.push(i128::from_le_bytes(lanes));
     }
 
     fn visit_v128_const(&mut self, value: wasmparser::V128) -> Self::Output {
-        self.instructions.push(Instruction::V128Const(self.v128_constants.len() as u32));
-        self.v128_constants.push(value.i128());
+        self.instructions.push(Instruction::V128Const(self.data.v128_constants.len() as u32));
+        self.data.v128_constants.push(value.i128());
     }
 }
 
@@ -593,7 +593,7 @@ impl<R: WasmModuleResources> FunctionBuilder<R> {
             validator,
             local_addr_map,
             instructions: Vec::with_capacity(instr_capacity),
-            v128_constants: Vec::new(),
+            data: FunctionDataBuilder::default(),
             ctx_stack: Vec::with_capacity(256),
             errors: Vec::new(),
         }
diff --git a/crates/tinywasm/src/interpreter/executor.rs b/crates/tinywasm/src/interpreter/executor.rs
index 0091013..8316421 100644
--- a/crates/tinywasm/src/interpreter/executor.rs
+++ b/crates/tinywasm/src/interpreter/executor.rs
@@ -154,8 +154,7 @@ impl<'store, const BUDGETED: bool> Executor<'store, BUDGETED> {
                     let k = *keep as usize;
                     self.store.stack.values.stack_ref.truncate_keep(b as usize, k);
                 }
-                BranchTable(default_ip, len) => { self.exec_branch_table(*default_ip, *len); continue; }
-                BranchTableTarget {..} => {},
+                BranchTable(default_ip, start, len) => { self.exec_branch_table(*default_ip, *start, *len); continue; }
                 Return => { if self.exec_return() { return Ok(Some(())); } continue; }
                 LocalGet32(local_index) => self.store.stack.values.push(self.store.stack.values.local_get::<Value32>(&self.cf, *local_index))?,
                 LocalGet64(local_index) => self.store.stack.values.push(self.store.stack.values.local_get::<Value64>(&self.cf, *local_index))?,
@@ -696,15 +695,10 @@ impl<'store, const BUDGETED: bool> Executor<'store, BUDGETED> {
     }
 
     #[inline(always)]
-    fn exec_branch_table(&mut self, default_ip: u32, len: u32) {
+    fn exec_branch_table(&mut self, default_ip: u32, start: u32, len: u32) {
         let idx = self.store.stack.values.pop::<i32>();
-        let start = self.cf.instr_ptr + 1;
-
         let target_ip = if idx >= 0 && (idx as u32) < len {
-            match self.func.instructions.0.get((start + idx as u32) as usize) {
-                Some(Instruction::BranchTableTarget(ip)) => *ip,
-                _ => default_ip,
-            }
+            self.func.data.branch_table_targets.get((start + idx as u32) as usize).copied().unwrap_or(default_ip)
         } else {
             default_ip
         };
diff --git a/crates/types/src/instructions.rs b/crates/types/src/instructions.rs
index dedebc5..78b52b8 100644
--- a/crates/types/src/instructions.rs
+++ b/crates/types/src/instructions.rs
@@ -77,8 +77,7 @@ pub enum Instruction {
     DropKeep64(u16, u16),
     DropKeep128(u16, u16),
     DropKeepRef(u16, u16),
-    BranchTable(u32, u32),       // (default_landing_pad_ip, target_count) - followed by BranchTableTarget entries
-    BranchTableTarget(u32),      // (landing_pad_ip)
+    BranchTable(u32, u32, u32),  // (default_landing_pad_ip, branch_table_start, target_count)
     Return,
     Call(FuncAddr),
     CallSelf,
@@ -164,8 +163,6 @@ pub enum Instruction {
     F32ConvertI32S, F32ConvertI32U, F32ConvertI64S, F32ConvertI64U, F32DemoteF64,
     F64ConvertI32S, F64ConvertI32U, F64ConvertI64S, F64ConvertI64U, F64PromoteF32,
 
-    // Reinterpretations are parser no-ops and intentionally omitted.
-
     // Saturating Float-to-Int Conversions
     I32TruncSatF32S, I32TruncSatF32U, I32TruncSatF64S, I32TruncSatF64U,
     I64TruncSatF32S, I64TruncSatF32U, I64TruncSatF64S, I64TruncSatF64U,
diff --git a/crates/types/src/lib.rs b/crates/types/src/lib.rs
index 2f6107a..a56dc9f 100644
--- a/crates/types/src/lib.rs
+++ b/crates/types/src/lib.rs
@@ -318,6 +318,7 @@ impl<'de, T: serde::Deserialize<'de>> serde::Deserialize<'de> for ArcSlice<T> {
 #[cfg_attr(feature = "archive", derive(serde::Serialize, serde::Deserialize))]
 pub struct WasmFunctionData {
     pub v128_constants: Box<[i128]>,
+    pub branch_table_targets: Box<[u32]>,
 }
 
 /// A WebAssembly Module Export