refactor(codegen): introduce external function in code section

codegen/src/code/func.rs

@@ -0,0 +1,35 @@
+//! External Function for the code section.
+use opcodes::ShangHai as OpCode;
+
+/// External function in code section.
+#[derive(PartialEq, Eq, Debug, Clone, Hash)]
+pub struct ExtFunc {
+    /// Stack input.
+    pub stack_out: u8,
+    /// Stack output.
+    pub stack_in: u8,
+    /// The bytecode of the external function.
+    pub bytecode: Vec<u8>,
+}
+
+impl ExtFunc {
+    /// Function select.
+    pub fn select() -> Self {
+        Self {
+            stack_in: 2,
+            stack_out: 1,
+            bytecode: [
+                OpCode::JUMPDEST,
+                OpCode::POP,
+                OpCode::PUSH1,
+                OpCode::Data(0x06),
+                OpCode::ADD,
+                OpCode::JUMP,
+            ]
+            .iter()
+            .copied()
+            .map(|op| op.into())
+            .collect(),
+        }
+    }
+}

codegen/src/jump/code.rs → codegen/src/code/mod.rs

@@ -1,14 +1,16 @@
 //! Table for the code section.
 
-use crate::Func;
+pub use func::ExtFunc;
 use indexmap::IndexMap;
 
+mod func;
+
 /// Code section for EVM.
 #[derive(Default, Debug)]
 pub struct Code {
     offset: usize,
     /// Function table.
-    funcs: IndexMap<Func, usize>,
+    funcs: IndexMap<ExtFunc, usize>,
 }
 
 impl Code {
@@ -21,7 +23,7 @@ impl Code {
     }
 
     /// Get the functions in the code section.
-    pub fn funcs(&self) -> Vec<Func> {
+    pub fn funcs(&self) -> Vec<ExtFunc> {
         self.funcs.keys().cloned().collect()
     }
 
@@ -34,12 +36,12 @@ impl Code {
     }
 
     /// Add a function to the code section.
-    pub fn try_add_func(&mut self, func: Func) {
+    pub fn try_add_func(&mut self, func: ExtFunc) {
         if self.funcs.contains_key(&func) {
             return;
         }
 
-        let bytecode = func.bytecode();
+        let bytecode = func.bytecode.clone();
         let len = bytecode.len();
         self.funcs.insert(func, self.offset);
         self.offset += len;
@@ -51,7 +53,7 @@ impl Code {
     }
 
     /// Get the offset of a function.
-    pub fn offset_of(&self, func: &Func) -> Option<u16> {
+    pub fn offset_of(&self, func: &ExtFunc) -> Option<u16> {
         self.funcs.get(func).and_then(|i| (*i).try_into().ok())
     }
 
@@ -60,7 +62,7 @@ impl Code {
         let mut code = Vec::new();
         for func in self.funcs.keys() {
             tracing::debug!("add function to code section: {:?}", func);
-            code.extend(func.bytecode());
+            code.extend(func.bytecode.clone());
         }
         code
     }

codegen/src/import.rs

@@ -5,45 +5,12 @@ use std::{
 };
 
 use crate::{Error, MacroAssembler, Result};
-use opcodes::ShangHai as OpCode;
-
-/// Selects one of its first two operands based on whether
-/// its third operand is zero or not.
-const SELECT: [OpCode; 6] = [
-    OpCode::JUMPDEST,
-    OpCode::POP,
-    OpCode::PUSH1,
-    OpCode::Data(0x06),
-    OpCode::ADD,
-    OpCode::JUMP,
-];
-
-/// Function `sload` from EVM which is not available in WASM.
-const SLOAD: [OpCode; 7] = [
-    OpCode::JUMPDEST,
-    OpCode::SLOAD,
-    OpCode::SWAP1,
-    OpCode::PUSH1,
-    OpCode::Data(0x05),
-    OpCode::ADD,
-    OpCode::JUMP,
-];
-
-/// Function `sstore` from EVM which is not available in WASM.
-const SSTORE: [OpCode; 6] = [
-    OpCode::JUMPDEST,
-    OpCode::SSTORE,
-    OpCode::PUSH1,
-    OpCode::Data(0x05),
-    OpCode::ADD,
-    OpCode::JUMP,
-];
 
 /// EVM built-in function.
 #[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq, Hash)]
 pub enum Func {
-    /// Run function select.
-    Select,
+    // EVM functions.
+    //
     /// Run function sload.
     Sload,
     /// Run function sstore.
@@ -59,7 +26,9 @@ pub enum Func {
     /// Run function log4.
     Log4,
 
+    //
     // Zinkc helper functions
+    //
     /// Emit ABI to the compiler.
     EmitABI,
 }
@@ -68,7 +37,6 @@ impl Func {
     /// Stack input size.
     pub fn stack_in(&self) -> u8 {
         match self {
-            Self::Select => 3,
             Self::Sload => 1,
             Self::Sstore => 2,
             Self::Log0 => 2,
@@ -83,7 +51,6 @@ impl Func {
     /// Stack output size.
     pub fn stack_out(&self) -> u8 {
         match self {
-            Self::Select => 1,
             Self::Sload => 1,
             Self::Sstore => 0,
             Self::Log0 => 0,
@@ -98,38 +65,10 @@ impl Func {
     /// Pre-processing for the function.
     pub fn prelude(&self, masm: &mut MacroAssembler) -> Result<()> {
         match self {
-            Self::Select => {
-                masm._pc()?;
-                masm._swap2()?;
-                masm._swap1()
-            }
             Self::Sstore => masm._swap1(),
             _ => Ok(()),
         }
     }
-
-    /// Get the bytecode of the function.
-    pub fn bytecode(&self) -> Vec<u8> {
-        match self {
-            Self::Select => SELECT.to_vec(),
-            Self::Sload => SLOAD.to_vec(),
-            Self::Sstore => SSTORE.to_vec(),
-            _ => unimplemented!("not implemented for {:?}", self),
-        }
-        .into_iter()
-        .map(|op| op.into())
-        .collect()
-    }
-
-    /// If this function should be embedded in
-    /// the main code.
-    ///
-    /// TODO: return `false` for functions that
-    /// are necessary to just stay in the code
-    /// section #109
-    pub fn is_embedded(&self) -> bool {
-        true
-    }
 }
 
 impl TryFrom<(&str, &str)> for Func {

codegen/src/jump/mod.rs

@@ -1,15 +1,14 @@
 //! Jump table implementation.
 
-pub use self::{code::Code, table::JumpTable};
-use crate::Func;
+use crate::code::ExtFunc;
+pub use table::JumpTable;
 
-mod code;
 mod pc;
 mod relocate;
 mod table;
 
 /// Jump types
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub enum Jump {
     /// offset to the program counter.
     Offset(u16),
@@ -19,7 +18,7 @@ pub enum Jump {
     /// Jump to function.
     Func(u32),
     /// External function.
-    ExtFunc(Func),
+    ExtFunc(ExtFunc),
 }
 
 impl Jump {

codegen/src/jump/pc.rs

@@ -74,7 +74,7 @@ impl JumpTable {
                     }
                 }
 
-                Ok((k, *v))
+                Ok((k, v.clone()))
             })
             .collect::<Result<_>>()?;
 

codegen/src/jump/table.rs

@@ -1,9 +1,6 @@
 //! Jump Table
 
-use crate::{
-    jump::{Code, Jump},
-    Error, Func, Result,
-};
+use crate::{code::ExtFunc, jump::Jump, Code, Error, Result};
 use std::collections::BTreeMap;
 
 /// Jump table implementation.
@@ -33,15 +30,15 @@ impl JumpTable {
         Ok(())
     }
 
-    /// Register program counter for code section.
+    /// Register the start of the program counter
+    /// of the code section.
     pub fn code_offset(&mut self, offset: u16) {
         self.code.shift(offset);
     }
 
     /// Register a external function.
-    pub fn ext(&mut self, pc: u16, func: Func) {
-        tracing::debug!("register external function: {:?}", func);
-        self.code.try_add_func(func);
+    pub fn ext(&mut self, pc: u16, func: ExtFunc) {
+        self.code.try_add_func(func.clone());
         self.jump.insert(pc, Jump::ExtFunc(func));
     }
 
@@ -89,7 +86,7 @@ impl JumpTable {
             Jump::Offset(offset) => Ok(*offset),
             Jump::Label(label) => Ok(*label),
             Jump::Func(func) => Ok(*self.func.get(func).ok_or(Error::FuncNotFound(*func))?),
-            Jump::ExtFunc(ext) => Ok(self.code.offset_of(ext).ok_or(Error::ExtNotFound(*ext))?),
+            Jump::ExtFunc(ext) => Ok(self.code.offset_of(ext).ok_or(Error::ExtFuncNotFound)?),
         }
     }
 }

codegen/src/lib.rs

@@ -5,13 +5,14 @@
 pub use crate::{
     abi::{ToLSBytes, Type},
     asm::Assembler,
+    code::Code,
     codegen::CodeGen,
     control::{ControlStack, ControlStackFrame, ControlStackFrameType},
     data::DataSet,
     dispatcher::{Dispatcher, Function, Functions},
     export::Exports,
     import::{Func, Imports},
-    jump::{Code, JumpTable},
+    jump::JumpTable,
     local::{LocalSlot, Locals},
     masm::MacroAssembler,
     result::{Error, Result},
@@ -21,6 +22,7 @@ use smallvec::SmallVec;
 pub mod abi;
 mod asm;
 mod backtrace;
+mod code;
 mod codegen;
 mod control;
 mod data;

codegen/src/result.rs

@@ -25,8 +25,8 @@ pub enum Error {
     #[error("Program counter {0} already exists in jump table")]
     DuplicateJump(u16),
     /// Failed to find ext function index in jump table.
-    #[error("External function {0:?} not found in jump table")]
-    ExtNotFound(crate::Func),
+    #[error("External function not found in jump table")]
+    ExtFuncNotFound,
     /// Failed to find function index in jump table.
     #[error("Function {0} not found in jump table")]
     FuncNotFound(u32),

codegen/src/visitor/call.rs

@@ -57,17 +57,7 @@ impl CodeGen {
             .ok_or(Error::ImportedFuncNotFound(index))?;
 
         func.prelude(&mut self.masm)?;
-        if func.is_embedded() {
-            tracing::debug!("embed imported function {func:?} at index {index}");
-            self.call_embedded(func)
-        } else {
-            tracing::debug!("call imported function {func:?} at index {index}");
-            self.call_external(func)
-        }
-    }
 
-    /// Call embedded functions
-    fn call_embedded(&mut self, func: Func) -> Result<()> {
         match func {
             Func::Sstore => self.masm._sstore(),
             Func::Sload => self.masm._sload(),
@@ -82,19 +72,4 @@ impl CodeGen {
             }
         }
     }
-
-    /// Call external functions
-    pub fn call_external(&mut self, func: Func) -> Result<()> {
-        // record the current program counter and
-        // pass it to the callee function.
-        self.masm._pc()?;
-
-        // register function to the jump table.
-        self.table.ext(self.masm.pc_offset(), func);
-        self.masm.decrement_sp(func.stack_in())?;
-        self.masm._jump()?;
-        self.masm._jumpdest()?;
-        self.masm.increment_sp(func.stack_out())?;
-        Ok(())
-    }
 }

codegen/src/visitor/control.rs

@@ -1,8 +1,9 @@
 //! Control flow visitors
 
 use crate::{
+    code::ExtFunc,
     control::{ControlStackFrame, ControlStackFrameType},
-    CodeGen, Func, Result,
+    CodeGen, Result,
 };
 use wasmparser::{BlockType, BrTable};
 
@@ -90,16 +91,20 @@ impl CodeGen {
     /// STACK: [val1, val2, cond] -> \[val1\] if cond is non-zero, \[val2\] otherwise.
     pub fn _select(&mut self) -> Result<()> {
         tracing::trace!("select");
-        let func = Func::Select;
-        func.prelude(&mut self.masm)?;
-        self.masm.decrement_sp(func.stack_in())?;
+        let ext = ExtFunc::select();
 
-        // This if for pushing the PC of jumpdest.
-        self.masm.increment_sp(1)?;
-        self.table.ext(self.masm.pc_offset(), Func::Select);
+        // Reorder the stack inputs.
+        self.masm._pc()?;
+        self.masm._swap2()?;
+        self.masm._swap1()?;
+
+        // Decrement the stack pointer.
+        self.masm.decrement_sp(ext.stack_in)?;
+
+        self.table.ext(self.masm.pc_offset(), ext.clone());
         self.masm._jumpi()?;
         self.masm._jumpdest()?;
-        self.masm.increment_sp(func.stack_out())?;
+        self.masm.increment_sp(ext.stack_out)?;
         Ok(())
     }