From 80ded1fa9eb26b8b90227247d1e14fe0a711bf0a Mon Sep 17 00:00:00 2001
From: clearloop <26088946+clearloop@users.noreply.github.com>
Date: Thu, 3 Aug 2023 17:04:25 +0800
Subject: [PATCH] fix(codegen): initial offset of locals (#61)

* refactor(codegen): use dup in local.get

* fix(codegen): merge of function tables

* fix(codegen): swap target in local.set

* chore(codegen): test local.tee

* refactor(locals): initial offset of locals

* fix(codegen): incorrect stack pointer increment in function call

* feat(example): fibonacci

* chore(zint): decrease the preset gas limit for evm
---
 Cargo.lock                           |  1 +
 codegen/src/asm.rs                   |  3 +++
 codegen/src/codegen.rs               | 31 ++++++++++++++------------
 codegen/src/control.rs               | 16 +++++++++-----
 codegen/src/jump/mod.rs              |  2 +-
 codegen/src/jump/pc.rs               |  7 +++---
 codegen/src/jump/relocate.rs         |  7 +-----
 codegen/src/local.rs                 | 26 ++--------------------
 codegen/src/masm/cmp.rs              | 14 ++++++++++++
 codegen/src/masm/mod.rs              | 15 ++++++++-----
 codegen/src/result.rs                |  3 +++
 codegen/src/visitor/control.rs       | 31 ++++++++++++++------------
 codegen/src/visitor/handlers.rs      | 33 ++++++++++++++++++++++++++--
 codegen/src/visitor/local.rs         | 27 +++++++++++++----------
 codegen/src/visitor/mod.rs           |  2 +-
 codegen/src/visitor/system.rs        |  5 ++---
 compiler/tests/add.rs                |  9 ++++++++
 compiler/tests/call.rs               | 29 +++++++++++++++++++-----
 compiler/tests/loop.rs               | 15 ++++++++++++-
 compiler/tests/recursion.rs          | 28 ++++++++++++++++-------
 compiler/wat/call/as_if.wat          | 17 ++++++++++++++
 compiler/wat/i32add/tee.wat          | 15 +++++++++++++
 compiler/wat/loop/as_br_if.wat       | 12 ++++++++++
 compiler/wat/recursion/fibonacci.wat | 28 +++++++++++------------
 zint/Cargo.toml                      |  1 +
 zint/src/evm.rs                      |  5 ++++-
 26 files changed, 260 insertions(+), 122 deletions(-)
 create mode 100644 compiler/wat/call/as_if.wat
 create mode 100644 compiler/wat/i32add/tee.wat
 create mode 100644 compiler/wat/loop/as_br_if.wat

diff --git a/Cargo.lock b/Cargo.lock
index 1b306cfa0..4fa8c830a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1716,4 +1716,5 @@ dependencies = [
  "hex",
  "revm-interpreter",
  "revm-primitives",
+ "tracing",
 ]
diff --git a/codegen/src/asm.rs b/codegen/src/asm.rs
index 37b07d5a2..d9ec82b9e 100644
--- a/codegen/src/asm.rs
+++ b/codegen/src/asm.rs
@@ -47,6 +47,7 @@ impl Assembler {
             return Err(Error::StackOverflow(self.sp));
         }
 
+        // tracing::debug!("increment sp: {items} -> {}", self.sp);
         Ok(())
     }
 
@@ -56,6 +57,8 @@ impl Assembler {
             .sp
             .checked_sub(items)
             .ok_or(Error::StackUnderflow(self.sp, items))?;
+
+        // tracing::debug!("decrement sp: {items} -> {}", self.sp);
         Ok(())
     }
 
diff --git a/codegen/src/codegen.rs b/codegen/src/codegen.rs
index ac80a54cf..b48fdb7b5 100644
--- a/codegen/src/codegen.rs
+++ b/codegen/src/codegen.rs
@@ -1,12 +1,11 @@
 //! Code generation implementation.
 use crate::{
-    abi::Type,
     control::ControlStack,
     jump::JumpTable,
     local::{LocalSlot, LocalSlotType, Locals},
     masm::MacroAssembler,
     validator::ValidateThenVisit,
-    Buffer, Result,
+    Buffer, Error, Result,
 };
 use wasmparser::{FuncType, FuncValidator, LocalsReader, OperatorsReader, ValidatorResources};
 
@@ -49,7 +48,7 @@ impl CodeGen {
         if !is_main {
             // Mock the stack frame for the callee function
             //
-            // STACK: PC
+            // STACK: PC + params
             codegen.masm.increment_sp(1 + params_count)?;
             codegen.masm._jumpdest()?;
             codegen.masm.shift_pc(params_count, true)?;
@@ -70,30 +69,29 @@ impl CodeGen {
         locals: &mut LocalsReader<'_>,
         validator: &mut FuncValidator<ValidatorResources>,
     ) -> Result<()> {
-        // Define locals in function parameters.
-        for (idx, param) in self.env.params().iter().enumerate() {
-            let sp = if self.is_main { None } else { Some(idx + 1) };
+        let mut sp = if self.is_main { 0 } else { 1 };
 
+        // Define locals in function parameters.
+        for param in self.env.params() {
             self.locals
                 .push(LocalSlot::new(*param, LocalSlotType::Parameter, sp));
+            sp += 1;
         }
 
         // Define locals in function body.
         //
         // Record the offset for validation.
-        let mut pc = self.env.params().len();
         while let Ok((count, val)) = locals.read() {
-            let sp = {
-                self.masm.increment_sp(1)?;
-                pc += 1;
-                Some(pc)
-            };
-
             let validation_offset = locals.original_position();
             for _ in 0..count {
+                // Init locals with zero.
+                self.masm.push(&[0])?;
+
+                // Define locals.
                 self.locals
                     .push(LocalSlot::new(val, LocalSlotType::Variable, sp));
-                self.masm.increment_mp(val.align())?;
+
+                sp += 1;
             }
 
             validator.define_locals(validation_offset, count, val)?;
@@ -120,6 +118,11 @@ impl CodeGen {
 
     /// Finish code generation.
     pub fn finish(self, jump_table: &mut JumpTable, pc: u16) -> Result<Buffer> {
+        let sp = self.masm.sp();
+        if !self.is_main && self.masm.sp() != self.env.results().len() as u8 {
+            return Err(Error::StackNotBalanced(sp));
+        }
+
         jump_table.merge(self.table, pc)?;
         Ok(self.masm.buffer().into())
     }
diff --git a/codegen/src/control.rs b/codegen/src/control.rs
index c10cda7b1..1c165c4ec 100644
--- a/codegen/src/control.rs
+++ b/codegen/src/control.rs
@@ -36,14 +36,23 @@ pub struct ControlStackFrame {
     ///
     /// Could be useful for validation.
     result: BlockType,
+
+    /// Original stack pointer.
+    pub original_sp: u8,
 }
 
 impl ControlStackFrame {
     /// Create a new control stack frame.
-    pub fn new(ty: ControlStackFrameType, original_pc_offset: u16, result: BlockType) -> Self {
+    pub fn new(
+        ty: ControlStackFrameType,
+        original_pc_offset: u16,
+        original_sp: u8,
+        result: BlockType,
+    ) -> Self {
         Self {
             ty,
             original_pc_offset,
+            original_sp,
             result,
         }
     }
@@ -57,11 +66,6 @@ impl ControlStackFrame {
     pub fn result(&self) -> BlockType {
         self.result
     }
-
-    // /// Check if the control stack frame is an if block with else.
-    // pub fn if_with_else(&self) -> bool {
-    //     self.ty == ControlStackFrameType::If(true)
-    // }
 }
 
 /// The control stack.
diff --git a/codegen/src/jump/mod.rs b/codegen/src/jump/mod.rs
index d167f3c9e..291c10700 100644
--- a/codegen/src/jump/mod.rs
+++ b/codegen/src/jump/mod.rs
@@ -59,7 +59,7 @@ impl JumpTable {
     /// counter of the target jump table.
     pub fn merge(&mut self, mut table: Self, pc: u16) -> Result<()> {
         if pc != 0 {
-            table.shift_pc(pc, pc)?;
+            table.shift_pc(0, pc)?;
         }
 
         for (pc, jump) in table.jump.into_iter() {
diff --git a/codegen/src/jump/pc.rs b/codegen/src/jump/pc.rs
index d7dbfb788..2c59a461c 100644
--- a/codegen/src/jump/pc.rs
+++ b/codegen/src/jump/pc.rs
@@ -8,8 +8,9 @@ use crate::{
 impl JumpTable {
     /// Shift program counter for all items.
     pub fn shift_pc(&mut self, start: u16, offset: u16) -> Result<()> {
-        tracing::debug!("shift pc: {} -> {}", start, offset);
+        tracing::debug!("shift pc: 0x{:x} -> 0x{:x}", start, offset);
         self.shift_label_pc(start, offset)?;
+        self.shift_label_target(start, offset)?;
         self.shift_func_target(start, offset)?;
 
         Ok(())
@@ -44,7 +45,7 @@ impl JumpTable {
     pub fn shift_func_target(&mut self, start: u16, offset: u16) -> Result<()> {
         self.func.values_mut().try_for_each(|v| {
             if *v > start {
-                tracing::debug!("shift function target: {} -> {}", v, *v + offset);
+                tracing::debug!("shift function target: 0x{:x} -> 0x{:x}", v, *v + offset);
                 *v += offset;
                 if *v > BUFFER_LIMIT as u16 {
                     return Err(Error::InvalidPC(*v as usize));
@@ -65,7 +66,7 @@ impl JumpTable {
             .map(|(k, v)| {
                 let mut k = *k;
                 if k > start {
-                    tracing::debug!("shift label pc: {} -> {}", k, k + offset);
+                    tracing::debug!("shift label pc: 0x{:x} -> 0x{:x}", k, k + offset);
                     k += offset;
                     if k > BUFFER_LIMIT as u16 {
                         return Err(Error::InvalidPC(k as usize));
diff --git a/codegen/src/jump/relocate.rs b/codegen/src/jump/relocate.rs
index 684169903..37ea327af 100644
--- a/codegen/src/jump/relocate.rs
+++ b/codegen/src/jump/relocate.rs
@@ -68,13 +68,8 @@ fn pc(buffer: &mut Buffer, original_pc: u16, target_pc: u16, offset: u16) -> Res
         new_buffer.push(OpCode::PUSH2.into());
     }
 
-    tracing::debug!(
-        "run pc relocation: 0x{:x} -> 0x{:x}",
-        original_pc,
-        target_pc
-    );
     let pc_offset = target_pc.to_ls_bytes();
-    tracing::debug!("push bytes: {:x?} at {}", pc_offset, original_pc);
+    tracing::debug!("push bytes: 0x{:x?} at 0x{:x}", pc_offset, original_pc);
     new_buffer.extend_from_slice(&pc_offset);
     new_buffer.extend_from_slice(&rest_buffer);
 
diff --git a/codegen/src/local.rs b/codegen/src/local.rs
index 4d443c9e0..1e1052ce2 100644
--- a/codegen/src/local.rs
+++ b/codegen/src/local.rs
@@ -28,12 +28,12 @@ pub struct LocalSlot {
     ty: LocalSlotType,
 
     /// Stack pointer of the local slot.
-    pub sp: Option<usize>,
+    pub sp: usize,
 }
 
 impl LocalSlot {
     /// Create a new local slot.
-    pub fn new(inner: ValType, ty: LocalSlotType, sp: Option<usize>) -> Self {
+    pub fn new(inner: ValType, ty: LocalSlotType, sp: usize) -> Self {
         Self { inner, ty, sp }
     }
 
@@ -104,28 +104,6 @@ impl Locals {
         self.inner.push(slot.into())
     }
 
-    /// Shift local stack pointers.
-    pub fn shift_sp(&mut self, index: usize, sp: usize) -> Result<()> {
-        let local = self.get_mut(index)?;
-        let src_sp = local.sp.ok_or_else(|| Error::InvalidLocalIndex(index))?;
-        local.sp = Some(sp);
-
-        for (idx, local) in self.inner.iter_mut().enumerate() {
-            if idx == index {
-                continue;
-            }
-
-            if let Some(local_sp) = local.sp {
-                if local_sp > src_sp {
-                    // TODO: Artihmetic checks
-                    local.sp = Some(local_sp - 1);
-                }
-            }
-        }
-
-        Ok(())
-    }
-
     /// Get the length of locals
     pub fn len(&self) -> usize {
         self.inner.len()
diff --git a/codegen/src/masm/cmp.rs b/codegen/src/masm/cmp.rs
index e57b116d4..65f16f7b9 100644
--- a/codegen/src/masm/cmp.rs
+++ b/codegen/src/masm/cmp.rs
@@ -4,6 +4,20 @@ use crate::{MacroAssembler, Result};
 use opcodes::ShangHai as OpCode;
 
 impl MacroAssembler {
+    /// Greater than or equal comparison.
+    ///
+    /// TODO: refactor this.
+    pub fn _ge(&mut self) -> Result<()> {
+        self._sgt()
+    }
+
+    /// Greater than comparison.
+    ///
+    /// Using lt due to order of stack.
+    pub fn _gt(&mut self) -> Result<()> {
+        self.asm._lt()
+    }
+
     /// Sign-agnostic compare unequal.
     pub fn _ne(&mut self) -> Result<()> {
         self.emit_op(OpCode::EQ)?;
diff --git a/codegen/src/masm/mod.rs b/codegen/src/masm/mod.rs
index fec2e4cf9..bc3805100 100644
--- a/codegen/src/masm/mod.rs
+++ b/codegen/src/masm/mod.rs
@@ -66,7 +66,7 @@ impl MacroAssembler {
 
     /// Place n bytes on stack.
     pub fn push(&mut self, bytes: &[u8]) -> Result<()> {
-        tracing::trace!("push bytes: {:?}", bytes);
+        tracing::trace!("push bytes: 0x{:x?}", bytes);
         let len = bytes.len();
         match len {
             0 => self.asm._push0(),
@@ -174,11 +174,13 @@ impl MacroAssembler {
     /// parameters. This is used by the callee function for jumping
     /// back to the caller function.
     pub fn shift_pc(&mut self, count: u8, from_top: bool) -> Result<()> {
-        self.swap(count)?;
+        let mut swaps = 0;
 
         if from_top {
-            if count > 1 {
-                return self.shift_pc(count - 1, from_top);
+            swaps = count;
+            while swaps > 0 {
+                self.swap(swaps)?;
+                swaps -= 1;
             }
         } else {
             // TODO: Optimize the shift logic when params lt 2.
@@ -188,8 +190,9 @@ impl MacroAssembler {
             // in total.
             //
             // if count > 2 {}
-            if count > 0 {
-                return self.shift_pc(count - 1, from_top);
+            while swaps < count {
+                swaps += 1;
+                self.swap(swaps)?;
             }
         }
 
diff --git a/codegen/src/result.rs b/codegen/src/result.rs
index 4242bc38f..c581d6278 100644
--- a/codegen/src/result.rs
+++ b/codegen/src/result.rs
@@ -57,6 +57,9 @@ pub enum Error {
     /// Failed to decrement stack pointer.
     #[error("Stack underflow, current stack items {0}, expect at least {1}")]
     StackUnderflow(u8, u8),
+    /// Failed to pop stack.
+    #[error("Stack not balanced, current stack items {0}")]
+    StackNotBalanced(u8),
 }
 
 /// Codegen result
diff --git a/codegen/src/visitor/control.rs b/codegen/src/visitor/control.rs
index 41b2c4814..5907362cf 100644
--- a/codegen/src/visitor/control.rs
+++ b/codegen/src/visitor/control.rs
@@ -13,6 +13,7 @@ impl CodeGen {
         let frame = ControlStackFrame::new(
             ControlStackFrameType::If(false),
             self.masm.pc_offset(),
+            self.masm.sp(),
             blockty,
         );
         self.control.push(frame);
@@ -29,8 +30,12 @@ impl CodeGen {
     /// The begeinning of a block construct. A sequence of
     /// instructions with a label at the end.
     pub fn _block(&mut self, blockty: BlockType) -> Result<()> {
-        let frame =
-            ControlStackFrame::new(ControlStackFrameType::Block, self.masm.pc_offset(), blockty);
+        let frame = ControlStackFrame::new(
+            ControlStackFrameType::Block,
+            self.masm.pc_offset(),
+            self.masm.sp(),
+            blockty,
+        );
         self.masm._jumpdest()?;
         self.control.push(frame);
 
@@ -40,10 +45,16 @@ impl CodeGen {
     /// A block with a label which may be used to
     /// form loops.
     pub fn _loop(&mut self, blockty: BlockType) -> Result<()> {
-        let frame =
-            ControlStackFrame::new(ControlStackFrameType::Loop, self.masm.pc_offset(), blockty);
+        let frame = ControlStackFrame::new(
+            ControlStackFrameType::Loop,
+            self.masm.pc_offset(),
+            self.masm.sp(),
+            blockty,
+        );
 
+        self.masm._jumpdest()?;
         self.control.push(frame);
+
         Ok(())
     }
 
@@ -55,6 +66,7 @@ impl CodeGen {
         let frame = ControlStackFrame::new(
             ControlStackFrameType::Else,
             self.masm.pc_offset(),
+            self.masm.sp(),
             last_frame.result(),
         );
         self.control.push(frame);
@@ -111,17 +123,8 @@ impl CodeGen {
     /// - End of function.
     /// - End of program.
     pub fn _end(&mut self) -> Result<()> {
-        // If inside an if frame, pop the frame and patch
-        // the program counter.
         if let Ok(frame) = self.control.pop() {
-            match frame.ty {
-                ControlStackFrameType::If(true) => {
-                    // TODO: fix this for nested if-else.
-                    self.handle_return()
-                }
-                ControlStackFrameType::Block => self.masm._jumpdest(),
-                _ => self.handle_jumpdest(frame.original_pc_offset),
-            }
+            self.handle_frame_popping(frame)
         } else if !self.is_main {
             self.handle_call_return()
         } else {
diff --git a/codegen/src/visitor/handlers.rs b/codegen/src/visitor/handlers.rs
index 78a1faa49..f07a76d23 100644
--- a/codegen/src/visitor/handlers.rs
+++ b/codegen/src/visitor/handlers.rs
@@ -1,6 +1,6 @@
 //! Case handlers
 
-use crate::{CodeGen, Error, Result, ToLSBytes};
+use crate::{CodeGen, ControlStackFrame, ControlStackFrameType, Error, Result, ToLSBytes};
 
 impl CodeGen {
     pub(crate) fn handle_empty_return(&mut self) -> Result<()> {
@@ -35,8 +35,22 @@ impl CodeGen {
         let results = self.env.results();
         tracing::debug!("handle call return: {:?}", results);
 
+        let len = results.len() as u8;
+        let sp = self.masm.sp();
+        for i in 0..len {
+            // TODO: arthmetic overflow.
+            //
+            // 2 is for PC & self.
+            self.masm.swap(sp - i - 2)?;
+        }
+
+        tracing::debug!("cleaning frame stack, target: {}", len + 1);
+        while self.masm.sp() > len + 1 {
+            self.masm._drop()?;
+        }
+
         // TODO: handle the length of results > u8::MAX.
-        self.masm.shift_pc(results.len() as u8, false)?;
+        self.masm.shift_pc(len, false)?;
         self.masm.push(&[0x04])?;
         self.masm._add()?;
         self.masm._jump()?;
@@ -44,6 +58,21 @@ impl CodeGen {
         Ok(())
     }
 
+    /// Handle the popping of a frame.
+    ///
+    /// TODO: validate stack IO for all frames (#59)
+    pub(crate) fn handle_frame_popping(&mut self, frame: ControlStackFrame) -> Result<()> {
+        match frame.ty {
+            ControlStackFrameType::If(true) => {
+                // TODO: fix this for nested if-else.
+                self.handle_return()
+            }
+            ControlStackFrameType::Block => self.masm._jumpdest(),
+            ControlStackFrameType::Loop => Ok(()),
+            _ => self.handle_jumpdest(frame.original_pc_offset),
+        }
+    }
+
     /// Handle jumpdest.
     pub(crate) fn handle_jumpdest(&mut self, original_pc: u16) -> Result<()> {
         self.table.label(original_pc, self.masm.pc_offset());
diff --git a/codegen/src/visitor/local.rs b/codegen/src/visitor/local.rs
index bd94b336a..d5f34cd2f 100644
--- a/codegen/src/visitor/local.rs
+++ b/codegen/src/visitor/local.rs
@@ -14,17 +14,23 @@ impl CodeGen {
     }
 
     /// This instruction sets the value of a variable.
-    pub fn _local_set(&mut self, index: u32) -> Result<()> {
-        let index = index as usize;
+    pub fn _local_set(&mut self, local_index: u32) -> Result<()> {
+        let index = local_index as usize;
+        let sp = self.masm.sp();
+        let local = self.locals.get(index)?;
+        let local_sp = local.sp as u8;
+
+        tracing::debug!("local_set: {index} {local_sp} {sp}");
+        self.masm.swap(sp - local_sp - 1)?;
+        self.masm._drop()?;
 
-        // Override the stack pointer of the local
-        self.locals.get_mut(index)?.sp = Some(self.masm.sp() as usize);
         Ok(())
     }
 
     /// This _local_tee is like _local_set, but it also returns the value
-    /// on the stack, however, in our implementation, they are the same.
+    /// on the stack.
     pub fn _local_tee(&mut self, index: u32) -> Result<()> {
+        self.masm._dup1()?;
         self._local_set(index)?;
         Ok(())
     }
@@ -55,15 +61,12 @@ impl CodeGen {
         }
 
         let local = self.locals.get(local_index)?;
-        let local_sp = local.sp.ok_or(Error::LocalNotOnStack(local_index))? as u8;
+        let local_sp = local.sp as u8;
         let sp = self.masm.sp();
 
-        if local_sp == sp {
-            return Ok(());
-        }
-
+        tracing::debug!("local_get: {local_index} {local_sp} {sp}");
         // TODO: Arthmetic checks
-        self.masm.swap(sp - local_sp)?;
-        self.locals.shift_sp(local_index, sp as usize)
+        self.masm.dup(sp - local_sp)?;
+        Ok(())
     }
 }
diff --git a/codegen/src/visitor/mod.rs b/codegen/src/visitor/mod.rs
index 424cb730d..83133db15 100644
--- a/codegen/src/visitor/mod.rs
+++ b/codegen/src/visitor/mod.rs
@@ -201,7 +201,7 @@ impl<'a> VisitOperator<'a> for CodeGen {
         ],
         signed_and_float: [add, sub, mul, eq, ne],
         map: {
-            all: [ge => sgt, le => slt],
+            all: [ge => ge, le => slt],
             integer: [rem => mod],
         },
         mem: {
diff --git a/codegen/src/visitor/system.rs b/codegen/src/visitor/system.rs
index 7e79b5dfe..147e6555b 100644
--- a/codegen/src/visitor/system.rs
+++ b/codegen/src/visitor/system.rs
@@ -23,10 +23,9 @@ impl CodeGen {
         // register the call index to the jump table.
         self.table.call(self.masm.pc_offset(), function_index)?;
 
-        // mock the stack output of the counter
+        // jump to the callee function
         //
-        // the program counter operators should be relocated afterwards.
-        self.masm.asm.increment_sp(1)?;
+        // TODO: check the stack output.
         self.masm._jump()?;
         self.masm._jumpdest()?;
         Ok(())
diff --git a/compiler/tests/add.rs b/compiler/tests/add.rs
index fcd82424c..50bc0fde3 100644
--- a/compiler/tests/add.rs
+++ b/compiler/tests/add.rs
@@ -26,3 +26,12 @@ fn locals() -> Result<()> {
     assert_eq!(info.ret, [30.to_bytes32()].concat());
     Ok(())
 }
+
+#[test]
+fn tee() -> Result<()> {
+    let bytecode = common::load("i32add", "tee")?;
+    let info = EVM::run(&bytecode, &[]);
+
+    assert_eq!(info.ret, [30.to_bytes32()].concat());
+    Ok(())
+}
diff --git a/compiler/tests/call.rs b/compiler/tests/call.rs
index 390c44a85..d6a0fc292 100644
--- a/compiler/tests/call.rs
+++ b/compiler/tests/call.rs
@@ -6,24 +6,41 @@ use zint::{Bytes32, InstructionResult, EVM};
 
 mod common;
 
+#[test]
+fn dummy() -> Result<()> {
+    let bytecode = common::load("call", "dummy")?;
+    let info = EVM::run(&bytecode, &[]);
+
+    assert_eq!(info.instr, InstructionResult::Return);
+    assert_eq!(info.ret, []);
+
+    Ok(())
+}
+
 #[test]
 fn params() -> Result<()> {
     let bytecode = common::load("call", "params")?;
 
     let input = [1.to_bytes32(), 2.to_bytes32()].concat();
     let info = EVM::run(&bytecode, &input);
-    tracing::trace!("info: {:?}", info);
 
     assert_eq!(info.ret, 3.to_bytes32());
     Ok(())
 }
 
 #[test]
-fn dummy() -> Result<()> {
-    let bytecode = common::load("call", "dummy")?;
-    let info = EVM::run(&bytecode, &[]);
+fn as_if() -> Result<()> {
+    let bytecode = common::load("call", "as_if")?;
+    let info = EVM::run(&bytecode, &0.to_bytes32());
+    assert_eq!(info.ret, 0.to_bytes32());
 
-    assert_eq!(info.instr, InstructionResult::Return);
-    assert_eq!(info.ret, []);
+    let info = EVM::run(&bytecode, &1.to_bytes32());
+    assert_eq!(info.ret, 1.to_bytes32());
+
+    let info = EVM::run(&bytecode, &2.to_bytes32());
+    assert_eq!(info.ret, 41.to_bytes32());
+
+    let info = EVM::run(&bytecode, &3.to_bytes32());
+    assert_eq!(info.ret, 42.to_bytes32());
     Ok(())
 }
diff --git a/compiler/tests/loop.rs b/compiler/tests/loop.rs
index 77b7b9d13..14dbcf4d4 100644
--- a/compiler/tests/loop.rs
+++ b/compiler/tests/loop.rs
@@ -1,7 +1,7 @@
 //! loop tests
 
 use anyhow::Result;
-use zint::{Bytes32, EVM};
+use zint::{Bytes32, InstructionResult, EVM};
 
 mod common;
 
@@ -14,3 +14,16 @@ fn singular() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn as_br_if() -> Result<()> {
+    let bytecode = common::load("loop", "as_br_if")?;
+
+    let info = EVM::run(&bytecode, 0.to_bytes32().as_ref());
+    assert_eq!(info.ret, 7.to_bytes32());
+
+    let info = EVM::run(&bytecode, 1.to_bytes32().as_ref());
+    assert_eq!(info.instr, InstructionResult::OutOfGas);
+
+    Ok(())
+}
diff --git a/compiler/tests/recursion.rs b/compiler/tests/recursion.rs
index 0b39d0144..ec29b7756 100644
--- a/compiler/tests/recursion.rs
+++ b/compiler/tests/recursion.rs
@@ -3,7 +3,7 @@ use zint::{Bytes32, EVM};
 
 mod common;
 
-#[ignore]
+// #[ignore]
 #[test]
 fn fibonacci() -> Result<()> {
     let bytecode = common::load("recursion", "fibonacci")?;
@@ -12,13 +12,25 @@ fn fibonacci() -> Result<()> {
     let info = EVM::run(&bytecode, &0.to_bytes32());
     assert_eq!(0.to_bytes32().to_vec(), info.ret);
 
-    // // x = 1
-    // let info = EVM::run(&bytecode, &1.to_bytes32());
-    // assert_eq!(1.to_bytes32().to_vec(), info.ret);
-    //
-    // // x = 3
-    // let info = EVM::run(&bytecode, &3.to_bytes32());
-    // assert_eq!(3.to_bytes32().to_vec(), info.ret);
+    // x = 1
+    let info = EVM::run(&bytecode, &1.to_bytes32());
+    assert_eq!(1.to_bytes32().to_vec(), info.ret);
+
+    // x = 2
+    let info = EVM::run(&bytecode, &2.to_bytes32());
+    assert_eq!(1.to_bytes32().to_vec(), info.ret);
+
+    // x = 3
+    let info = EVM::run(&bytecode, &3.to_bytes32());
+    assert_eq!(2.to_bytes32().to_vec(), info.ret);
+
+    // x = 4
+    let info = EVM::run(&bytecode, &4.to_bytes32());
+    assert_eq!(3.to_bytes32().to_vec(), info.ret);
+
+    // x = 5
+    let info = EVM::run(&bytecode, &5.to_bytes32());
+    assert_eq!(5.to_bytes32().to_vec(), info.ret);
 
     Ok(())
 }
diff --git a/compiler/wat/call/as_if.wat b/compiler/wat/call/as_if.wat
new file mode 100644
index 000000000..65f2d7177
--- /dev/null
+++ b/compiler/wat/call/as_if.wat
@@ -0,0 +1,17 @@
+(module
+  (type (;0;) (func (param i32) (result i32)))
+  (func (;0;) (type 0) (param i32) (result i32)
+    local.get 0
+    call 1)
+  (func (;1;) (type 0) (param i32) (result i32)
+    (local i32)
+    local.get 0
+    i32.const 2
+    i32.ge_u
+    if  ;; label = @1
+      i32.const 39
+      local.set 1
+    end
+    local.get 0
+    local.get 1
+    i32.add))
diff --git a/compiler/wat/i32add/tee.wat b/compiler/wat/i32add/tee.wat
new file mode 100644
index 000000000..f4aa981a2
--- /dev/null
+++ b/compiler/wat/i32add/tee.wat
@@ -0,0 +1,15 @@
+(module
+    (func (result i32)
+        (local $foo i32)
+        (local $bar i32)
+
+        (i32.const 10)
+        (local.tee $foo)
+
+        (i32.const 20)
+        (local.set $bar)
+
+        (local.get $bar)
+        i32.add
+    )
+)
diff --git a/compiler/wat/loop/as_br_if.wat b/compiler/wat/loop/as_br_if.wat
new file mode 100644
index 000000000..7de3eed23
--- /dev/null
+++ b/compiler/wat/loop/as_br_if.wat
@@ -0,0 +1,12 @@
+(module
+  (func (export "br_if") (param i32) (result i32)
+    loop ;; label = @1
+      i32.const 0
+      local.get 0
+      i32.gt_s
+      br_if 0 (;1;)
+    end
+
+    i32.const 7
+  )
+)
diff --git a/compiler/wat/recursion/fibonacci.wat b/compiler/wat/recursion/fibonacci.wat
index 4ae547563..f95146463 100644
--- a/compiler/wat/recursion/fibonacci.wat
+++ b/compiler/wat/recursion/fibonacci.wat
@@ -10,20 +10,20 @@
     i32.ge_u
     if  ;; label = @1
       loop  ;; label = @2
-        local.get 0
-        i32.const 1
-        i32.sub
-        call 1
-        local.get 1
-        i32.add
-        local.set 1
-        local.get 0
-        i32.const 2
-        i32.sub
-        local.tee 0
-        i32.const 1
-        i32.gt_u
-        br_if 0 (;@2;)
+        local.get 0    ;; 1
+        i32.const 1    ;; 2
+        i32.sub        ;; 1
+        call 1         ;; 1
+        local.get 1    ;; 2
+        i32.add        ;; 1
+        local.set 1    ;; 0
+        local.get 0    ;; 1
+        i32.const 2    ;; 2
+        i32.sub        ;; 1
+        local.tee 0    ;; 1
+        i32.const 1    ;; 2
+        i32.gt_u       ;; 1
+        br_if 0 (;@2;) ;; 2 -> 0
       end
     end
     local.get 0
diff --git a/zint/Cargo.toml b/zint/Cargo.toml
index 19aa9c3de..5016cf3ac 100644
--- a/zint/Cargo.toml
+++ b/zint/Cargo.toml
@@ -13,3 +13,4 @@ documentation = "https://docs.rs/ztime"
 hex.workspace = true
 revm-interpreter.workspace = true
 revm-primitives.workspace = true
+tracing.workspace = true
diff --git a/zint/src/evm.rs b/zint/src/evm.rs
index 169eb3237..93761370a 100644
--- a/zint/src/evm.rs
+++ b/zint/src/evm.rs
@@ -59,6 +59,9 @@ impl EVM {
     /// Run a contract.
     pub fn run(btyecode: &[u8], input: &[u8]) -> Info {
         let mut evm = Self::new(btyecode, input);
-        evm.execute()
+        let info = evm.execute();
+        tracing::debug!("{info:?}");
+
+        info
     }
 }