refactor(codegen): get locals from memory in internal functions (#245)

* refactor(codegen): remove control handlers

* feat(codegen): check and patch return code at function end

* refactor(codegen): get params from memory in internal functions

* feat(codegen): label reserved memory bytes in compilation

* feat(codegen): allocate memory for locals

* feat(codegen): embed slots in env

* refactor(call): preprocess of function calls

* refactor(codegen): remove initial empty values for locals in internal functions

* feat(codegen): make slots based from the previous function

* fix(compiler): tests of br_if

* chore(codegen): ban recursion

codegen/src/asm.rs

@@ -46,11 +46,11 @@ impl Assembler {
             return Ok(());
         }
 
-        // tracing::trace!(
-        //     "increment stack pointer {}({items}) -> {}",
-        //     self.sp,
-        //     self.sp + items
-        // );
+        tracing::debug!(
+            "increment stack pointer {}.add({items}) -> {}",
+            self.sp,
+            self.sp + items
+        );
         self.sp += items;
 
         // TODO: fix this limitation: should be 1024. (#127)
@@ -67,15 +67,18 @@ impl Assembler {
             return Ok(());
         }
 
-        tracing::trace!(
-            "decrement stack pointer {}({items}) -> {}",
+        tracing::debug!(
+            "decrement stack pointer {}.sub({items}) -> {}",
             self.sp,
             self.sp - items
         );
-        self.sp = self
-            .sp
-            .checked_sub(items)
-            .ok_or(Error::StackUnderflow(self.sp, items))?;
+        self.sp = if self.sp == items {
+            0
+        } else {
+            self.sp
+                .checked_sub(items)
+                .ok_or(Error::StackUnderflow(self.sp, items))?
+        };
 
         Ok(())
     }

codegen/src/backtrace.rs

@@ -23,6 +23,11 @@ impl Backtrace {
         self.instrs.pop_last().unwrap_or_default().1
     }
 
+    /// Get the last byte
+    pub fn last(&self) -> Option<Vec<u8>> {
+        self.instrs.last_key_value().map(|(_, op)| op.clone())
+    }
+
     /// Pop the last `n` operands from the backtrace.
     pub fn popn(&mut self, n: usize) -> Vec<Vec<u8>> {
         let mut r: Vec<Vec<u8>> = Default::default();

codegen/src/codegen/function.rs

@@ -9,6 +9,7 @@ use crate::{
     wasm::Env,
     Buffer, Error, Result,
 };
+use opcodes::ShangHai as OpCode;
 use wasmparser::{FuncType, FuncValidator, LocalsReader, OperatorsReader, ValidatorResources};
 use zabi::Abi;
 
@@ -37,11 +38,6 @@ pub struct Function {
 impl Function {
     /// Create a new code generator.
     pub fn new(env: Env, ty: FuncType, abi: Option<Abi>, is_main: bool) -> Result<Self> {
-        let mut params_count = 0;
-        if !is_main {
-            params_count = ty.params().len() as u8;
-        }
-
         let is_external = abi.is_some();
         let mut codegen = Self {
             abi,
@@ -61,15 +57,16 @@ impl Function {
 
         // post process program counter and stack pointer.
         if is_external {
-            codegen.masm.increment_sp(1)?;
+            // codegen.masm.increment_sp(1)?;
+            tracing::debug!("<External function>");
             codegen.masm._jumpdest()?;
         } else {
             // Mock the stack frame for the callee function
             //
-            // STACK: PC + params
-            codegen.masm.increment_sp(1 + params_count)?;
+            // STACK: [ PC ]
+            tracing::debug!("<Internal function>");
+            codegen.masm.increment_sp(1)?;
             codegen.masm._jumpdest()?;
-            codegen.masm.shift_stack(params_count, true)?;
         }
 
         Ok(codegen)
@@ -80,7 +77,7 @@ impl Function {
     /// 1. the function parameters.
     /// 2. function body locals.
     ///
-    /// NOTE: we don't care about the origin offset of the locals.
+    /// NOTE: we don't care about the original offset of the locals.
     /// bcz we will serialize the locals to an index map anyway.
     pub fn emit_locals(
         &mut self,
@@ -100,10 +97,8 @@ impl Function {
         //
         // Record the offset for validation.
         while let Ok((count, val)) = locals.read() {
-            let validation_offset = locals.original_position();
             for _ in 0..count {
-                // Init locals with zero.
-                self.masm.push(&[0])?;
+                // TODO: the below here is outdated, sp is not required anymore after #245
 
                 // Define locals.
                 self.locals
@@ -112,6 +107,7 @@ impl Function {
                 sp += 1;
             }
 
+            let validation_offset = locals.original_position();
             validator.define_locals(validation_offset, count, val)?;
         }
 
@@ -131,6 +127,12 @@ impl Function {
             ops.visit_operator(&mut validate_then_visit)???;
         }
 
+        if (self.abi.is_some() || self.is_main)
+            && self.masm.buffer().last() != Some(&OpCode::RETURN.into())
+        {
+            self._end()?;
+        }
+
         Ok(())
     }
 

codegen/src/jump/table.rs

@@ -47,7 +47,7 @@ impl JumpTable {
         self.jump.insert(pc, Jump::Label(label));
     }
 
-    /// Register a label.
+    /// Register a label at the specific PC offset
     pub fn offset(&mut self, pc: u16, offset: u16) {
         self.jump.insert(pc, Jump::Offset(offset));
     }

codegen/src/masm/mod.rs

@@ -87,6 +87,14 @@ impl MacroAssembler {
     /// Place n bytes on stack.
     pub fn push(&mut self, bytes: &[u8]) -> Result<()> {
         tracing::trace!("push bytes: 0x{:x?}", bytes);
+
+        // TODO: support PUSH0 #247
+        //
+        // if !bytes.iter().any(|b| *b != 0) {
+        //     self.asm._push0()?;
+        //     return Ok(());
+        // }
+
         let len = bytes.len();
         match len {
             0 => self.asm._push0(),

codegen/src/masm/ret.rs

@@ -32,20 +32,15 @@ impl MacroAssembler {
     /// Handle the return of a call.
     pub fn call_return(&mut self, results: &[ValType]) -> Result<()> {
         let len = results.len() as u8;
-        let sp = self.sp();
-        tracing::trace!("current stack items: {sp}");
-        for i in 0..len {
-            // TODO: arthmetic overflow.
-            //
-            // 2 is for PC and self.
-            self.swap(sp.saturating_sub(i).saturating_sub(2))?;
-        }
 
         tracing::trace!("cleaning frame stack, target: {}", len + 1);
         while self.sp() > len + 1 {
             self._drop()?;
         }
 
+        // Shift stack to prompt the jump instruction,
+        // what about just dup it?
+        //
         // TODO: handle the length of results > u8::MAX.
         self.shift_stack(len, false)?;
         self._jump()

codegen/src/result.rs

@@ -3,6 +3,9 @@
 /// Codegen error
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
+    /// Any error
+    #[error("{0}")]
+    Anyhow(#[from] anyhow::Error),
     /// Failed to parse function ABI.
     #[error(transparent)]
     Abi(#[from] zabi::result::Error),

codegen/src/visitor/call.rs

@@ -1,6 +1,10 @@
 //! call instructions
 
-use crate::{wasm::HostFunc, Error, Function, Result};
+use crate::{
+    wasm::{HostFunc, ToLSBytes},
+    Error, Function, Result,
+};
+use anyhow::anyhow;
 use opcodes::ShangHai as OpCode;
 
 impl Function {
@@ -31,26 +35,60 @@ impl Function {
 
     /// Call internal functions
     fn call_internal(&mut self, index: u32) -> Result<()> {
+        if self.env.index == Some(index) {
+            return Err(anyhow!(
+                "Recursion is no more supported in this version, see https://github.com/zink-lang/zink/issues/248"
+            )
+            .into());
+        }
+
         tracing::trace!("call internal function: index={index}");
-        // record the current program counter and
-        // pass it to the callee function.
-        self.table.offset(self.masm.pc_offset(), 6);
+        let (params, results) = self.env.funcs.get(&index).unwrap_or(&(0, 0));
+
+        // TODO: adapat the case that the params is larger than 0xff (#247)
+        //
+        // 1. record the program counter of the end of this expression
+        // call and pass it to the callee function.
+        //
+        // [ ..,
+        //   <PC>,
+        //   params[SWAP], params[PUSH, SLOT, MSTORE],
+        //   PUSH, PC, JUMP, <JUMPDEST>
+        // ]
+        // <- selfparams[PUSH, OFFSET, CALLDATALOAD]
+        //
+        // 2. move PC before the params in stack
+        self.table
+            .offset(self.masm.pc_offset(), 5 + 4 * (*params as u16));
         self.masm.increment_sp(1)?;
-        self.masm._jumpdest()?;
 
-        // Call an internal function.
+        // Stack
+        // =====
         //
-        // register the call index to the jump table.
+        // from  [ <PARAMS>, PC ]
+        // to    [ PC, <PARAMS> ]
+        self.masm.shift_stack(*params as u8, true)?;
+
+        // Call an internal function.
         //
+        // 1. store params in memory
+        // 2. register the call index to the jump table.
+        let reserved = self.env.slots.get(&index).unwrap_or(&0);
+        for i in (0..*params).rev() {
+            tracing::trace!("storing local at {} for function {index}", i + reserved);
+            self.masm.push(&((i + reserved) * 0x20).to_ls_bytes())?;
+            self.masm._mstore()?;
+        }
+
         // TODO: support same pc different jumps. (#160)
         self.table.call(self.masm.pc_offset(), index);
 
         // jump to the callee function
-        //
-        // TODO: check the stack output.
         self.masm._jump()?;
         self.masm._jumpdest()?;
 
+        // Stack: [ , ..results ]
+        self.masm.increment_sp(*results as u8)?;
         Ok(())
     }
 
@@ -74,7 +112,7 @@ impl Function {
             HostFunc::Evm(OpCode::LOG3) => self.log(3),
             HostFunc::Evm(OpCode::LOG4) => self.log(4),
             HostFunc::Evm(op) => self.masm.emit_op(op),
-            HostFunc::NoOp => Ok(()),
+            HostFunc::NoOp | HostFunc::Label(_) => Ok(()),
             _ => {
                 tracing::error!("unsupported host function {func:?}");
                 Err(Error::UnsupportedHostFunc(func))

codegen/src/visitor/control.rs

@@ -130,20 +130,21 @@ impl Function {
 
     /// Handle the end of instructions for different situations.
     ///
-    /// TODO: (#28)
-    ///
     /// - End of control flow operators.
     /// - End of function.
     /// - End of program.
     pub fn _end(&mut self) -> Result<()> {
         if let Ok(frame) = self.control.pop() {
-            self.handle_frame_popping(frame)
-        } else if self.is_main || self.abi.is_some() {
+            return self.handle_frame_popping(frame);
+        }
+
+        let results = self.ty.results();
+        if self.is_main || self.abi.is_some() {
             tracing::trace!("end of main function");
-            self.handle_return()
+            self.masm.main_return(results)
         } else {
             tracing::trace!("end of call");
-            self.handle_call_return()
+            self.masm.call_return(results)
         }
     }
 
@@ -160,4 +161,25 @@ impl Function {
     pub fn _nop(&mut self) -> Result<()> {
         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) => Ok(()),
+            ControlStackFrameType::Block => self.masm._jumpdest(),
+            ControlStackFrameType::Loop => Ok(()),
+            _ => {
+                self.table
+                    .label(frame.original_pc_offset, self.masm.pc_offset());
+
+                // TODO: Check the stack output and make decisions
+                // how to handle the results.
+
+                // Emit JUMPDEST after at the end of the control flow.
+                self.masm._jumpdest()
+            }
+        }
+    }
 }