From c71e18d66d35bff4633c2a3af457785b7329ce82 Mon Sep 17 00:00:00 2001 From: Steve Fan <29133953+stevefan1999-personal@users.noreply.github.com> Date: Thu, 3 Sep 2026 09:34:12 +0800 Subject: [PATCH] fix(parser): do not fuse local.get/local.tee across a branch target `visit_local_tee` peeked at the last emitted instruction and, when it was a `local.get`, replaced the pair with `LocalCopy(src, dst); LocalGet(dst)`. It never checked whether the tee's position is a jump target. When the `local.get` is a block's fallthrough result and the `local.tee` is the first instruction after `end`, every branch to that label lands on `LocalGet(dst)` instead of the tee: the branch's value is left dangling on the stack and a stale local is pushed on top. Loop starts and if/else joins have the same shape. TinyCC compiled with clang -Os/-Oz hit this in musl's realloc: the `if (!p) { r = malloc(n); br 1 }` path returned n instead of r, so the caller wrote into address 1200 and later trapped with an out-of-bounds load, while V8 and wabt ran the same module correctly. The lowering step now always emits a plain `LocalTee`. The fusion moves into the peephole rewriter, which only matches within a basic block and so cannot cross a label. Regression test covers the block-end, if/else-join and loop-start cases. --- crates/parser/src/optimize/rewrite.rs | 9 +++ crates/parser/src/visit.rs | 33 +++----- .../tinywasm/tests/branch_target_local_tee.rs | 77 +++++++++++++++++++ 3 files changed, 96 insertions(+), 23 deletions(-) create mode 100644 crates/tinywasm/tests/branch_target_local_tee.rs diff --git a/crates/parser/src/optimize/rewrite.rs b/crates/parser/src/optimize/rewrite.rs index 4a64f77..366a094 100644 --- a/crates/parser/src/optimize/rewrite.rs +++ b/crates/parser/src/optimize/rewrite.rs @@ -859,6 +859,9 @@ fn rewrite_local_tee32( replace!(output, *read, 1 => Instruction::SubConstTee32(I32LocalArg { value, local: dst })); } Instruction::LocalGet32(src) if src == dst => replace!(output, *read, 1 => Instruction::LocalGet32(src)), + Instruction::LocalGet32(src) => { + replace!(output, *read, 1 => [Instruction::LocalCopy32(src, dst), Instruction::LocalGet32(dst)]); + } Instruction::BinOpLocalLocal32(op, left, right) => { let replacement = if op == BinOp::IAdd { Instruction::AddLocalLocalTee32(LocalTripleArg { left, right, dst }) @@ -931,6 +934,9 @@ fn rewrite_local_tee64( replace!(output, *read, 1 => Instruction::SubConstTee64(PackedOp::new(dst, packed.index))); } Instruction::LocalGet64(src) if src == dst => replace!(output, *read, 1 => Instruction::LocalGet64(src)), + Instruction::LocalGet64(src) => { + replace!(output, *read, 1 => [Instruction::LocalCopy64(src, dst), Instruction::LocalGet64(dst)]); + } Instruction::BinOpLocalLocal64(op, left, right) => { let index = data.push_operand64(Operand64::<(u16, u16, u16)>::new(left, right, dst))?; replace!(output, *read, 1 => Instruction::BinOpLocalLocalTee64(PackedOp::new(op, index))); @@ -967,6 +973,9 @@ fn rewrite_local_tee128( if *read > output.block_start { match output[*read - 1] { Instruction::LocalGet128(src) if src == dst => replace!(output, *read, 1 => Instruction::LocalGet128(src)), + Instruction::LocalGet128(src) => { + replace!(output, *read, 1 => [Instruction::LocalCopy128(src, dst), Instruction::LocalGet128(dst)]); + } Instruction::BinOpLocalLocal128(op, left, right) => { let index = data.push_operand64(Operand64::<(u16, u16, u16)>::new(left, right, dst))?; replace!(output, *read, 1 => Instruction::BinOpLocalLocalTee128(PackedOp::new(op, index))); diff --git a/crates/parser/src/visit.rs b/crates/parser/src/visit.rs index 724c1bf..7daef52 100644 --- a/crates/parser/src/visit.rs +++ b/crates/parser/src/visit.rs @@ -735,29 +735,16 @@ impl<'a> wasmparser::VisitOperator<'a> for FunctionBuilder<'_> { fn visit_local_tee(&mut self, idx: u32) -> Self::Output { let (size, local_idx) = self.local(idx)?; - self.apply_effect(&[size], &[size])?; - let src = match (size, self.instructions.last()) { - (ValueLane::S32, Some(Instruction::LocalGet32(src))) => Some(*src), - (ValueLane::S64, Some(Instruction::LocalGet64(src))) => Some(*src), - (ValueLane::S128, Some(Instruction::LocalGet128(src))) => Some(*src), - _ => None, - }; - if let Some(src) = src { - self.instructions.pop(); - let instructions = match size { - ValueLane::S32 => [Instruction::LocalCopy32(src, local_idx), Instruction::LocalGet32(local_idx)], - ValueLane::S64 => [Instruction::LocalCopy64(src, local_idx), Instruction::LocalGet64(local_idx)], - ValueLane::S128 => [Instruction::LocalCopy128(src, local_idx), Instruction::LocalGet128(local_idx)], - }; - self.instructions.extend(instructions); - } else { - self.instructions.push(size.select( - Instruction::LocalTee32(local_idx), - Instruction::LocalTee64(local_idx), - Instruction::LocalTee128(local_idx), - )); - } - Ok(()) + // No peephole here: this position may be a branch target (block end, + // loop start, if/else join), and fusing with the preceding `local.get` + // would move the label past the tee. The rewriter fuses the same pair + // within a basic block, where no branch can land between them. + let instruction = size.select( + Instruction::LocalTee32(local_idx), + Instruction::LocalTee64(local_idx), + Instruction::LocalTee128(local_idx), + ); + self.emit(&[size], &[size], instruction) } fn visit_block(&mut self, blockty: wasmparser::BlockType) -> Self::Output { diff --git a/crates/tinywasm/tests/branch_target_local_tee.rs b/crates/tinywasm/tests/branch_target_local_tee.rs new file mode 100644 index 0000000..a58a315 --- /dev/null +++ b/crates/tinywasm/tests/branch_target_local_tee.rs @@ -0,0 +1,77 @@ +use tinywasm::{ModuleInstance, Store}; + +/// A `local.get` right before a label and a `local.tee` right after it must +/// not be fused: branches land on the label with the value on the stack and +/// still have to execute the tee. +#[test] +fn branches_landing_on_local_tee_still_tee() -> Result<(), Box> { + let wasm = wat::parse_str( + r#" + (module + (func $forty_two (result i32) i32.const 42) + + ;; block end label: `br` out of the `if` carries the callee result + (func (export "block_end") (param i32 i32) (result i32) + block (result i32) + local.get 0 + i32.eqz + if + call $forty_two + br 1 + end + local.get 0 + end + local.tee 1 + drop + local.get 1) + + ;; if/else join label: the then-arm's implicit jump lands on the tee + (func (export "if_else_join") (param i32 i32) (result i32) + local.get 0 + if (result i32) + call $forty_two + else + local.get 0 + end + local.tee 1 + drop + local.get 1) + + ;; loop start label: `br 0` re-enters at the tee with the loop param + (func (export "loop_start") (param i32 i32) (result i32) + local.get 0 + loop (param i32) (result i32) + local.tee 1 + i32.const 40 + i32.lt_u + if (result i32) + local.get 1 + i32.const 21 + i32.add + br 1 + else + local.get 1 + end + end) + ) + "#, + )?; + + let module = tinywasm::parse_bytes(&wasm)?; + let mut store = Store::default(); + let instance = ModuleInstance::instantiate(&mut store, &module, None)?; + + let block_end = instance.func::<(i32, i32), i32>(&store, "block_end")?; + assert_eq!(block_end.call(&mut store, (5, 7))?, 5, "fallthrough keeps local 0"); + assert_eq!(block_end.call(&mut store, (0, 7))?, 42, "branch path must tee the callee result"); + + let if_else_join = instance.func::<(i32, i32), i32>(&store, "if_else_join")?; + assert_eq!(if_else_join.call(&mut store, (0, 7))?, 0, "else arm keeps local 0"); + assert_eq!(if_else_join.call(&mut store, (1, 7))?, 42, "then arm must tee the callee result"); + + let loop_start = instance.func::<(i32, i32), i32>(&store, "loop_start")?; + assert_eq!(loop_start.call(&mut store, (40, 7))?, 40, "no iteration keeps the param"); + assert_eq!(loop_start.call(&mut store, (0, 7))?, 42, "re-entering the loop must tee the carried value"); + + Ok(()) +}