feat(codegen): introduce storage related built-in functions (#97)

* feat(codegen): introduce module func for built-in functions

* feat(codegen): introduce built-in function sload and sstore

* feat(compiler): passing imported zink functions to codegen

* feat(codegen): process base function index from imports

* feat(codegen): adjust imported functions in calls

* feat(compiler): add storage to the outputs of EVM tests

* feat(compiler): test sstore

* feat(codegen): correct the code section generation order

Cargo.toml

@@ -28,8 +28,7 @@ etc = "0.1.16"
 hex = "0.4.3"
 indexmap = "2.0.0"
 paste = "1.0.13"
-revm-interpreter = "1.1.2"
-revm-primitives = "1.1.2"
+revm = "3.3.0"
 semver = "1.0.18"
 serde = "1.0.171"
 smallvec = "1.11.0"

codegen/Cargo.toml

@@ -11,9 +11,10 @@ repository.workspace = true
 
 [dependencies]
 anyhow.workspace = true
-opcodes.workspace = true
 smallvec.workspace = true
 thiserror.workspace = true
 tracing.workspace = true
-wasmparser.workspace = true
+opcodes.workspace = true
 paste.workspace = true
+wasmparser.workspace = true
+indexmap.workspace = true

codegen/src/codegen.rs

@@ -5,7 +5,7 @@ use crate::{
     local::{LocalSlot, LocalSlotType, Locals},
     masm::MacroAssembler,
     validator::ValidateThenVisit,
-    Buffer, Error, Result,
+    Buffer, Error, Imports, Result,
 };
 use wasmparser::{FuncType, FuncValidator, LocalsReader, OperatorsReader, ValidatorResources};
 
@@ -23,13 +23,15 @@ pub struct CodeGen {
     pub(crate) masm: MacroAssembler,
     /// The jump table.
     pub(crate) table: JumpTable,
+    /// The imported functions.
+    pub(crate) imports: Imports,
     /// If this function is the main function.
     pub(crate) is_main: bool,
 }
 
 impl CodeGen {
     /// Create a new code generator.
-    pub fn new(env: FuncType, is_main: bool) -> Result<Self> {
+    pub fn new(env: FuncType, imports: Imports, is_main: bool) -> Result<Self> {
         let mut params_count = 0;
         if !is_main {
             params_count = env.params().len() as u8;
@@ -41,6 +43,7 @@ impl CodeGen {
             locals: Default::default(),
             masm: Default::default(),
             table: Default::default(),
+            imports,
             is_main,
         };
 

codegen/src/func.rs

@@ -0,0 +1,70 @@
+//! Built-in functions for EVM
+use opcodes::ShangHai as OpCode;
+
+/// Function `select` from WASM which is not avaiable in EVM.
+const SELECT: [OpCode; 6] = [
+    OpCode::JUMPDEST,
+    OpCode::POP,
+    OpCode::PUSH1,
+    OpCode::Data(0x06),
+    OpCode::ADD,
+    OpCode::JUMP,
+];
+
+/// Function `sload` from EVM which is not avaiable in WASM.
+const SLOAD: [OpCode; 7] = [
+    OpCode::JUMPDEST,
+    OpCode::SLOAD,
+    OpCode::SWAP1,
+    OpCode::PUSH1,
+    OpCode::Data(0x05),
+    OpCode::ADD,
+    OpCode::JUMP,
+];
+
+/// Function `sload` from EVM which is not avaiable in WASM.
+const SSTORE: [OpCode; 6] = [
+    OpCode::JUMPDEST,
+    OpCode::SSTORE,
+    OpCode::PUSH1,
+    OpCode::Data(0x05),
+    OpCode::ADD,
+    OpCode::JUMP,
+];
+
+/// Function selector.
+#[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq, Hash)]
+pub enum Func {
+    /// Run function select.
+    Select,
+    /// Run function sload.
+    Sload,
+    /// Run function sstore.
+    Sstore,
+}
+
+impl Func {
+    /// 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(),
+        }
+        .into_iter()
+        .map(|op| op.into())
+        .collect()
+    }
+}
+
+impl TryFrom<(&str, &str)> for Func {
+    type Error = ();
+
+    fn try_from(import: (&str, &str)) -> Result<Self, Self::Error> {
+        match import {
+            ("zink", "sload") => Ok(Self::Sload),
+            ("zink", "sstore") => Ok(Self::Sstore),
+            _ => Err(()),
+        }
+    }
+}

codegen/src/jump/code.rs

@@ -1,48 +1,22 @@
 //! Table for the code section.
 
-use opcodes::ShangHai as OpCode;
-use std::collections::HashMap;
-
-/// Code in code section.
-#[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq, Hash)]
-pub enum Func {
-    /// Run select.
-    Select,
-}
-
-impl Func {
-    /// Get the bytecode of the function.
-    pub fn bytecode(&self) -> Vec<u8> {
-        match self {
-            Self::Select => [
-                OpCode::JUMPDEST,
-                OpCode::POP,
-                OpCode::PUSH1,
-                OpCode::Data(0x06),
-                OpCode::ADD,
-                OpCode::JUMP,
-            ],
-        }
-        .into_iter()
-        .map(|op| op.into())
-        .collect()
-    }
-}
+use crate::func::Func;
+use indexmap::IndexMap;
 
 /// Code section for EVM.
 #[derive(Default, Debug)]
 pub struct Code {
     offset: usize,
     /// Function table.
-    funcs: HashMap<Func, usize>,
+    funcs: IndexMap<Func, usize>,
 }
 
 impl Code {
     /// Create a new code section.
     pub fn new() -> Self {
         Self {
             offset: 0,
-            funcs: HashMap::new(),
+            funcs: Default::default(),
         }
     }
 
@@ -85,6 +59,7 @@ impl Code {
     pub fn finish(&self) -> Vec<u8> {
         let mut code = Vec::new();
         for func in self.funcs.keys() {
+            tracing::debug!("add function to code section: {:?}", func);
             code.extend(func.bytecode());
         }
         code

codegen/src/jump/mod.rs

@@ -1,7 +1,7 @@
 //! Jump table implementation.
 
-use crate::{Error, Result};
-pub use code::{Code, Func};
+use crate::{Error, Func, Result};
+pub use code::Code;
 use std::collections::BTreeMap;
 
 mod code;
@@ -60,6 +60,7 @@ impl JumpTable {
 
     /// 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);
         self.jump.insert(pc, Jump::ExtFunc(func));
     }

codegen/src/lib.rs

@@ -7,17 +7,20 @@ pub use crate::{
     asm::Assembler,
     codegen::CodeGen,
     control::{ControlStack, ControlStackFrame, ControlStackFrameType},
+    func::Func,
     jump::{Code, JumpTable},
     local::{LocalSlot, Locals},
     masm::MacroAssembler,
     result::{Error, Result},
 };
+// use indexmap::IndexMap;
 use smallvec::SmallVec;
 
 pub mod abi;
 mod asm;
 mod codegen;
 mod control;
+mod func;
 mod jump;
 mod local;
 mod masm;
@@ -30,3 +33,7 @@ pub const BUFFER_LIMIT: usize = 0x6000;
 
 /// Code generation buffer.
 pub type Buffer = SmallVec<[u8; BUFFER_LIMIT]>;
+
+/// Imported functions.
+/// pub type Imports = IndexMap<u32, Func>;
+pub type Imports = Vec<Func>;

codegen/src/result.rs

@@ -18,12 +18,15 @@ pub enum Error {
     /// Failed to merge jump table.
     #[error("Program counter {0} already exists in jump table")]
     DuplicateJump(u16),
+    /// Failed to find imported function by index in jump table.
+    #[error("Imported Function {0} not found in jump table")]
+    ImportedFuncNotFound(u32),
     /// Failed to find function index in jump table.
     #[error("Function {0} not found in jump table")]
     FuncNotFound(u32),
     /// Failed to find ext function index in jump table.
     #[error("External function {0:?} not found in jump table")]
-    ExtNotFound(crate::jump::Func),
+    ExtNotFound(crate::Func),
     /// Failed to mark else block for if block.
     #[error("Invalid else block for if block at {0}")]
     InvalidElseBlock(u16),

codegen/src/visitor/control.rs

@@ -2,8 +2,7 @@
 
 use crate::{
     control::{ControlStackFrame, ControlStackFrameType},
-    jump::Func,
-    CodeGen, Result,
+    CodeGen, Func, Result,
 };
 use wasmparser::{BlockType, BrTable};
 
@@ -141,8 +140,10 @@ impl CodeGen {
         if let Ok(frame) = self.control.pop() {
             self.handle_frame_popping(frame)
         } else if !self.is_main {
+            tracing::debug!("end of call");
             self.handle_call_return()
         } else {
+            tracing::debug!("end of main function");
             self.handle_return()
         }
     }

codegen/src/visitor/handlers.rs

@@ -13,8 +13,9 @@ impl CodeGen {
 
     /// Handle the end of the function.
     pub(crate) fn handle_return(&mut self) -> Result<()> {
-        tracing::debug!("handle return");
         let results = self.env.results();
+        tracing::debug!("handle return, results: {results:?}");
+
         if results.is_empty() {
             return self.handle_empty_return();
         }

codegen/src/visitor/system.rs

@@ -1,6 +1,6 @@
 //! System instructions
 
-use crate::{CodeGen, Result};
+use crate::{CodeGen, Error, Func, Result};
 
 impl CodeGen {
     /// The call indirect instruction calls a function indirectly
@@ -15,17 +15,56 @@ impl CodeGen {
     }
 
     /// The call instruction calls a function specified by its index.
-    pub fn _call(&mut self, function_index: u32) -> Result<()> {
+    pub fn _call(&mut self, index: u32) -> Result<()> {
         // record the current program counter and
         // pass it to the callee function.
         self.masm._pc()?;
 
+        // TODO: check the safty of the function index.
+        let base = self.imports.len() as u32;
+
+        if base > index {
+            self.call_imported(index)
+        } else {
+            self.call_internal(base + index)
+        }
+    }
+
+    /// Call internal functions
+    fn call_internal(&mut self, index: u32) -> Result<()> {
+        // Call an internal function.
+        //
         // register the call index to the jump table.
-        self.table.call(self.masm.pc_offset(), function_index);
+        self.table.call(self.masm.pc_offset(), index);
 
         // jump to the callee function
         //
         // TODO: check the stack output.
+        self.masm._jump()?;
+        self.masm._jumpdest()?;
+
+        Ok(())
+    }
+
+    /// Call imported functions
+    fn call_imported(&mut self, index: u32) -> Result<()> {
+        // call an imported function.
+        //
+        // register the imported function index to the jump table.
+        let func = *self
+            .imports
+            .get(index as usize)
+            .ok_or(Error::ImportedFuncNotFound(index))?;
+
+        tracing::debug!("call imported function {func:?} at index {index}");
+        match func {
+            Func::Sstore => self.masm._swap2()?,
+            Func::Sload => self.masm._swap1()?,
+            _ => {}
+        }
+
+        self.table.ext(self.masm.pc_offset(), func);
+
         self.masm._jump()?;
         self.masm._jumpdest()?;
         Ok(())