feat(codegen): move return handlers to macro assembler

codegen/src/dispatcher.rs

@@ -5,7 +5,9 @@ use std::{
     collections::BTreeMap,
     ops::{Deref, DerefMut},
 };
-use wasmparser::{FuncValidator, FunctionBody, Operator, ValidatorResources};
+use wasmparser::{
+    FuncType, FuncValidator, FunctionBody, Operator, ValidatorResources, WasmModuleResources,
+};
 use zabi::Abi;
 
 /// Function with validator.
@@ -16,6 +18,26 @@ pub struct Function<'f> {
     pub body: FunctionBody<'f>,
 }
 
+impl Function<'_> {
+    /// Get function index.
+    pub fn index(&self) -> u32 {
+        self.validator.index()
+    }
+
+    /// Get the function signature.
+    pub fn sig(&self) -> Result<FuncType> {
+        let func_index = self.validator.index();
+        let sig = self
+            .validator
+            .resources()
+            .type_of_function(func_index)
+            .ok_or(Error::InvalidFunctionSignature)?
+            .clone();
+
+        Ok(sig)
+    }
+}
+
 /// Functions with indexes.
 #[derive(Default)]
 pub struct Functions<'f>(BTreeMap<u32, Function<'f>>);
@@ -146,19 +168,19 @@ impl Dispatcher {
     /// Emit selector to buffer.
     fn emit_selector(&mut self, selector: &Function<'_>, last: bool) -> Result<()> {
         let abi = self.load_abi(selector)?;
-        let selector = abi.selector();
+        let selector_bytes = abi.selector();
 
         tracing::debug!(
             "Emitting selector {:?} for function: {}",
-            selector,
+            selector_bytes,
             abi.name
         );
 
         let func = self.query_func(&abi.name)?;
 
-        // Jump to the end of the dispatcher.
+        // Jump to the end of the current function.
         //
-        // TODO: detect the bytes of the position.
+        // TODO: detect the bytes of the position. (#157)
         self.asm.increment_sp(1)?;
         let pc = self.asm.pc_offset();
         self.table.offset(pc, if last { 0xb } else { 0xc });
@@ -167,13 +189,14 @@ impl Dispatcher {
             self.asm._dup2()?;
         }
 
-        self.asm.push(&selector)?;
+        self.asm.push(&selector_bytes)?;
         self.asm._eq()?;
         self.asm.increment_sp(1)?;
 
         self.table.call(self.asm.pc_offset(), func);
         self.asm._jumpi()?;
         self.asm._jumpdest()?;
+        self.asm.main_return(selector.sig()?.results())?;
 
         Ok(())
     }

codegen/src/masm/integer.rs

@@ -12,14 +12,20 @@ impl MacroAssembler {
 
     /// Push a 32-bit integer value on the stack.
     pub fn _i32_const(&mut self, value: i32) -> Result<()> {
-        self.push(value.to_ls_bytes().as_ref())?;
-        Ok(())
+        if value == 0 {
+            self._push0()
+        } else {
+            self.push(value.to_ls_bytes().as_ref())
+        }
     }
 
     /// Push a 64-bit integer value on the stack.
     pub fn _i64_const(&mut self, value: i64) -> Result<()> {
-        self.push(value.to_ls_bytes().as_ref())?;
-        Ok(())
+        if value == 0 {
+            self._push0()
+        } else {
+            self.push(value.to_ls_bytes().as_ref())
+        }
     }
 
     /// Push a 32-bit float value on the stack.

codegen/src/masm/mod.rs

@@ -13,6 +13,7 @@ mod embed;
 mod float;
 mod integer;
 mod memory;
+mod ret;
 mod stack;
 
 /// EVM MacroAssembler.

codegen/src/masm/ret.rs

@@ -0,0 +1,58 @@
+//! Return handlers
+
+use crate::{Error, MacroAssembler, Result, ToLSBytes};
+use wasmparser::ValType;
+
+impl MacroAssembler {
+    /// Return with nothing.
+    pub(crate) fn handle_empty_return(&mut self) -> Result<()> {
+        self._push0()?;
+        self._push0()?;
+        self.asm._return()?;
+
+        Ok(())
+    }
+
+    /// Handle the end of the main function.
+    pub fn main_return(&mut self, results: &[ValType]) -> Result<()> {
+        if results.is_empty() {
+            return self.handle_empty_return();
+        }
+
+        let size = self.memory_write(results)?.size;
+        let offset =
+            self.mp_offset(|mp| mp.checked_sub(size).ok_or_else(|| Error::InvalidMP(0)))?;
+
+        self.push(&size.to_ls_bytes())?;
+        self.push(&offset)?;
+        self.asm._return()?;
+        Ok(())
+    }
+
+    /// 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();
+        for i in 0..len {
+            // TODO: arthmetic overflow.
+            //
+            // 2 is for PC & self.
+            self.swap(sp - i - 2)?;
+        }
+
+        tracing::debug!("cleaning frame stack, target: {}", len + 1);
+        while self.sp() > len + 1 {
+            self._drop()?;
+        }
+
+        // TODO: handle the length of results > u8::MAX.
+        self.shift_pc(len, false)?;
+
+        // TODO: [PC, PUSHN, BYTES, JUMPDEST]
+        self.push(&[0x04])?;
+        self._add()?;
+        self._jump()?;
+
+        Ok(())
+    }
+}

codegen/src/result.rs

@@ -42,6 +42,9 @@ pub enum Error {
     /// Failed to mark else block for if block.
     #[error("Invalid else block for if block at {0}")]
     InvalidElseBlock(u16),
+    /// Failed parse function signature.
+    #[error("Invalid function signature")]
+    InvalidFunctionSignature,
     /// Failed to get local with given index.
     #[error("Invalid local index {0}")]
     InvalidLocalIndex(usize),

codegen/src/visitor/handlers.rs

@@ -1,64 +1,22 @@
 //! Case handlers
 
-use crate::{CodeGen, ControlStackFrame, ControlStackFrameType, Error, Result, ToLSBytes};
+use crate::{CodeGen, ControlStackFrame, ControlStackFrameType, Result};
 
 impl CodeGen {
-    pub(crate) fn handle_empty_return(&mut self) -> Result<()> {
-        self.masm._push0()?;
-        self.masm._push0()?;
-        self.masm.asm._return()?;
-
-        Ok(())
-    }
-
     /// Handle the end of the function.
     pub(crate) fn handle_return(&mut self) -> Result<()> {
         let results = self.env.results();
         tracing::debug!("handle return, results: {results:?}");
 
-        if results.is_empty() {
-            return self.handle_empty_return();
-        }
-
-        let size = self.masm.memory_write(results)?.size;
-        let offset = self
-            .masm
-            .mp_offset(|mp| mp.checked_sub(size).ok_or_else(|| Error::InvalidMP(0)))?;
-
-        self.masm.push(&size.to_ls_bytes())?;
-        self.masm.push(&offset)?;
-        self.masm.asm._return()?;
-        Ok(())
+        self.masm.main_return(results)
     }
 
     /// Handle the return of a call.
     pub(crate) fn handle_call_return(&mut self) -> Result<()> {
         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(len, false)?;
-
-        // TODO: [PC, PUSHN, BYTES, JUMPDEST]
-        self.masm.push(&[0x04])?;
-        self.masm._add()?;
-        self.masm._jump()?;
-
-        Ok(())
+        self.masm.call_return(results)
     }
 
     /// Handle the popping of a frame.

compiler/src/compiler.rs

@@ -1,9 +1,9 @@
 //! Zink compiler
 
 use crate::{parser::Parser, Error, Result};
-use wasmparser::{FuncValidator, FunctionBody, ValidatorResources, WasmModuleResources};
 use zingen::{
-    Buffer, CodeGen, DataSet, Dispatcher, Exports, Functions, Imports, JumpTable, BUFFER_LIMIT,
+    Buffer, CodeGen, DataSet, Dispatcher, Exports, Function, Functions, Imports, JumpTable,
+    BUFFER_LIMIT,
 };
 
 /// Zink Compiler
@@ -39,7 +39,7 @@ impl Compiler {
         }
 
         for func in funcs.into_funcs() {
-            self.compile_func(data.clone(), imports.clone(), func.validator, func.body)?;
+            self.compile_func(data.clone(), imports.clone(), func)?;
         }
 
         self.finish()
@@ -71,15 +71,10 @@ impl Compiler {
         &mut self,
         dataset: DataSet,
         imports: Imports,
-        mut validator: FuncValidator<ValidatorResources>,
-        body: FunctionBody,
+        mut func: Function<'_>,
     ) -> Result<()> {
-        let func_index = validator.index();
-        let sig = validator
-            .resources()
-            .type_of_function(func_index)
-            .ok_or(Error::InvalidFunctionSignature)?
-            .clone();
+        let func_index = func.index();
+        let sig = func.sig()?;
 
         tracing::debug!("compile function {}: {:?}", func_index, sig);
 
@@ -90,11 +85,11 @@ impl Compiler {
         };
 
         let mut codegen = CodeGen::new(sig, dataset, imports, is_main)?;
-        let mut locals_reader = body.get_locals_reader()?;
-        let mut ops_reader = body.get_operators_reader()?;
+        let mut locals_reader = func.body.get_locals_reader()?;
+        let mut ops_reader = func.body.get_operators_reader()?;
 
-        codegen.emit_locals(&mut locals_reader, &mut validator)?;
-        codegen.emit_operators(&mut ops_reader, &mut validator)?;
+        codegen.emit_locals(&mut locals_reader, &mut func.validator)?;
+        codegen.emit_operators(&mut ops_reader, &mut func.validator)?;
 
         self.emit_buffer(func_index, codegen)?;
         Ok(())

compiler/src/result.rs

@@ -15,9 +15,6 @@ pub enum Error {
     /// Failed in code generation.
     #[error(transparent)]
     Codegen(#[from] zingen::Error),
-    /// Failed parse function signature.
-    #[error("Invalid function signature")]
-    InvalidFunctionSignature,
     /// Failed to parse WASM data with data reader.
     #[error("Invalid data offset")]
     InvalidDataOffset,