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

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(())
     }
 

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())
     }

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.

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() {

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

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

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

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

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

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

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 {