fix(parser): do not fuse local.get/local.tee across a branch target (#56)

I was trying to compile TCC to WASI and run it using tinywasm thanks to
the exception implementation (sigh, which is because of the use of
setjmp/longjmp for compiler rewind, that's the whole reason, WASI
syscall is really not that hard to say the least), but it went into a
mysterious bug that keeps trapping on the same address.

And yes, that means I did have a WASIp1 implementaion and I actually
implemented it with a VFS (nix the network socket just yet), so I can
theoretically have a cross-platform JIT compiler where I can have safe
reentrancy, and it is not until recently the TCC mob branches have their
TLS implementations, but it is all too late and they still rely on CRT.
It is also theoretically possible.

So my plan is to compile libtcc as small as possible to WASI using fat
LTO, and I keep getting trapped on the same address. Under NodeJS with a
shim WASI implementation, the whole libtcc fat LTO WASI actually works.
So it has to be a problem in tinywasm

Claude said this (it is Fable 5.1 btw):

> `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.
> 

I will try to use https://www.npmjs.com/package/@yowasp/clang to see if
there are more VM shenanigans, though a simple benchmark shows me that
the interpreter is 17x slower than native, and I guess the clang in
tinywasm wouldn't take less than 3 minutes to get a meaningful outcome
-- except if it blows up midway during execution.

Author: stevefan1999-personal

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)));

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 {

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<dyn core::error::Error>> {
+    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(())
+}