feat(codegen): load calldata for dispatcher

zink-lang · Oct 24, 2023 · 5d28e64 · 5d28e64
1 parent 884bd5f
commit 5d28e64
Show file tree

Hide file tree

Showing 16 changed files with 101 additions and 147 deletions.
diff --git a/codegen/src/abi.rs b/codegen/src/abi.rs
@@ -92,6 +92,7 @@ macro_rules! offset {
 
 offset! {
     (usize, 8),
+    (u64, 8),
     (i64, 8),
     (i32, 4),
     (u32, 4),

diff --git a/codegen/src/asm.rs b/codegen/src/asm.rs
@@ -97,7 +97,7 @@ impl Assembler {
     /// Mock the stack input and output for checking
     /// the stack usages.
     pub fn emit_op(&mut self, opcode: OpCode) -> Result<()> {
-        tracing::trace!("stack length: {:?}", self.sp);
+        // tracing::trace!("stack length: {:?}", self.sp);
         tracing::trace!("emit opcode: {:?}", opcode);
         self.decrement_sp(opcode.stack_in() as u8)?;
         self.emit(opcode.into());

diff --git a/codegen/src/codegen.rs b/codegen/src/codegen.rs
@@ -59,7 +59,7 @@ impl CodeGen {
             // STACK: PC + params
             codegen.masm.increment_sp(1 + params_count)?;
             codegen.masm._jumpdest()?;
-            codegen.masm.shift_pc(params_count, true)?;
+            codegen.masm.shift_stack(params_count, true)?;
         }
 
         Ok(codegen)

diff --git a/codegen/src/dispatcher.rs b/codegen/src/dispatcher.rs
@@ -1,6 +1,8 @@
 //! Code generator for EVM dispatcher.
 
-use crate::{code::ExtFunc, DataSet, Error, Exports, Imports, JumpTable, MacroAssembler, Result};
+use crate::{
+    code::ExtFunc, DataSet, Error, Exports, Imports, JumpTable, MacroAssembler, Result, ToLSBytes,
+};
 use std::{
     collections::BTreeMap,
     ops::{Deref, DerefMut},
@@ -179,15 +181,14 @@ impl<'d> Dispatcher<'d> {
     }
 
     /// Emit return of ext function.
-    fn ext_return(&mut self, func: u32) -> Result<()> {
-        let sig = self
-            .funcs
-            .get(&func)
-            .ok_or(Error::FuncNotFound(func))?
-            .sig()?;
-
+    fn ext_return(&mut self, sig: &FuncType) -> Result<()> {
+        self.asm.increment_sp(1)?;
         let asm = self.asm.clone();
-        self.asm.main_return(sig.results())?;
+
+        {
+            self.asm.main_return(sig.results())?;
+        }
+
         let bytecode = {
             let jumpdest = vec![0x5b];
             let ret = self.asm.buffer()[asm.buffer().len()..].to_vec();
@@ -204,6 +205,52 @@ impl<'d> Dispatcher<'d> {
         Ok(())
     }
 
+    // Process to the selected function.
+    //
+    // 1. drop selector.
+    // 2. load calldata to stack.
+    // 3. jump to the callee function.
+    fn process(&mut self, sig: &FuncType) -> Result<()> {
+        self.asm.increment_sp(1)?;
+        let asm = self.asm.clone();
+        let len = sig.params().len() as u8;
+
+        {
+            // TODO: check the safety of this.
+            //
+            // [ callee, ret, selector ] -> [ selector, callee, ret ]
+            self.asm.shift_stack(2, false)?;
+            // [ selector, callee, ret ] -> [ callee, ret ]
+            self.asm._drop()?;
+            // [ callee, ret ] -> [ param * len, callee, ret ]
+            for p in (0..len).rev() {
+                let offset = 4 + p * 32;
+                self.asm.push(&offset.to_ls_bytes())?;
+                self.asm._calldataload()?;
+            }
+
+            // [ param * len, callee, ret ] -> [ ret, param * len, callee ]
+            self.asm.shift_stack(len + 1, false)?;
+            // [ ret, param * len, callee ] -> [ callee, ret, param * len ]
+            self.asm.shift_stack(len + 1, false)?;
+            self.asm._jump()?;
+        }
+
+        let bytecode = {
+            let jumpdest = vec![0x5b];
+            let ret = self.asm.buffer()[asm.buffer().len()..].to_vec();
+            [jumpdest, ret].concat()
+        };
+        *self.asm = asm;
+        let ret = ExtFunc {
+            bytecode,
+            stack_in: len,
+            stack_out: 1,
+        };
+        self.table.ext(self.asm.pc_offset(), ret);
+        Ok(())
+    }
+
     /// Emit selector to buffer.
     fn emit_selector(&mut self, selector: &Function<'_>, last: bool) -> Result<()> {
         let abi = self.load_abi(selector)?;
@@ -216,27 +263,40 @@ impl<'d> Dispatcher<'d> {
         );
 
         let func = self.query_func(&abi.name)?;
+        let sig = self
+            .funcs
+            .get(&func)
+            .ok_or(Error::FuncNotFound(func))?
+            .sig()?;
 
         // Jump to the end of the current function.
         //
         // TODO: detect the bytes of the position. (#157)
-        self.asm.increment_sp(1)?;
-        self.ext_return(func)?;
+        self.ext_return(&sig)?;
+
+        // Prepare the `PC` of the callee function.
+        {
+            // TODO: remove this (#160)
+            self.asm._jumpdest()?;
+            self.asm.increment_sp(1)?;
+            self.table.call(self.asm.pc_offset(), func);
+        }
 
         if !last {
-            self.asm._dup2()?;
+            self.asm._dup3()?;
         } else {
-            self.asm._swap1()?;
+            self.asm._swap2()?;
         }
 
         self.asm.push(&selector_bytes)?;
         self.asm._eq()?;
-        self.asm.increment_sp(1)?;
-
-        self.table.call(self.asm.pc_offset(), func);
+        self.process(&sig)?;
         self.asm._jumpi()?;
 
         if !last {
+            // drop the PC of the previous callee function.
+            self.asm._drop()?;
+            // drop the PC of the previous callee function preprocessor.
             self.asm._drop()?;
         }
 

diff --git a/codegen/src/import.rs b/codegen/src/import.rs
@@ -4,7 +4,7 @@ use std::{
     ops::{Deref, DerefMut},
 };
 
-use crate::{Error, MacroAssembler, Result};
+use crate::{Error, Result};
 
 /// EVM built-in function.
 #[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq, Hash)]
@@ -61,14 +61,6 @@ impl Func {
             _ => 0,
         }
     }
-
-    /// Pre-processing for the function.
-    pub fn prelude(&self, masm: &mut MacroAssembler) -> Result<()> {
-        match self {
-            Self::Sstore => masm._swap1(),
-            _ => Ok(()),
-        }
-    }
 }
 
 impl TryFrom<(&str, &str)> for Func {

diff --git a/codegen/src/masm/mod.rs b/codegen/src/masm/mod.rs
@@ -193,7 +193,7 @@ impl MacroAssembler {
     /// Shift the program counter to the bottom or the top of the
     /// 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<()> {
+    pub fn shift_stack(&mut self, count: u8, from_top: bool) -> Result<()> {
         let mut swaps = 0;
 
         if from_top {

diff --git a/codegen/src/masm/ret.rs b/codegen/src/masm/ret.rs
@@ -46,7 +46,7 @@ impl MacroAssembler {
         }
 
         // TODO: handle the length of results > u8::MAX.
-        self.shift_pc(len, false)?;
+        self.shift_stack(len, false)?;
         self._jump()
     }
 }
diff --git a/codegen/src/visitor/call.rs b/codegen/src/visitor/call.rs
@@ -60,8 +60,6 @@ impl CodeGen {
             .get(&index)
             .ok_or(Error::ImportedFuncNotFound(index))?;
 
-        func.prelude(&mut self.masm)?;
-
         match func {
             Func::Sstore => self.masm._sstore(),
             Func::Sload => self.masm._sload(),

diff --git a/compiler/tests/storage.rs b/compiler/tests/storage.rs
@@ -45,19 +45,19 @@ fn basic() -> Result<()> {
 
 #[test]
 fn dispatcher() -> Result<()> {
-    let _bytecode = common::load_with_dispatcher("storage", "dispatcher")?;
+    let bytecode = common::load_with_dispatcher("storage", "dispatcher")?;
 
     // TODO: testing set (#122)
-    // {
-    //     let key = 0;
-    //     let value = 42;
-    //     let mut selector = zabi::selector(b"set(i32)").to_vec();
-    //     selector = [selector, 42.to_bytes32().to_vec()].concat();
-    //     let info = EVM::run(&bytecode, &selector);
-    //     assert_eq!(info.instr, InstructionResult::Return);
-    //     assert_eq!(info.ret, []);
-    //     assert_eq!(info.storage.get(&U256::from(key)), Some(&U256::from(value)));
-    // }
+    {
+        let key = 0;
+        let value = 42;
+        let mut selector = zabi::selector(b"set(i32)").to_vec();
+        selector = [selector, 42.to_bytes32().to_vec()].concat();
+        let info = EVM::run(&bytecode, &selector);
+        assert_eq!(info.instr, InstructionResult::Return);
+        assert_eq!(info.ret, []);
+        assert_eq!(info.storage.get(&U256::from(key)), Some(&U256::from(value)));
+    }
 
     // let info = EVM::run(&bytecode, &42.to_bytes32());
 

diff --git a/compiler/wat/storage/basic.wat b/compiler/wat/storage/basic.wat
@@ -4,8 +4,8 @@
   (import "evm" "sstore" (func (;0;) (type 1)))
   (import "evm" "sload" (func (;1;) (type 0)))
   (func (;2;) (type 0) (param i32) (result i32)
-    i32.const 0
     local.get 0
+    i32.const 0
     call 0
     i32.const 0
     call 1))
diff --git a/compiler/wat/storage/dispatcher.wat b/compiler/wat/storage/dispatcher.wat
@@ -9,8 +9,8 @@
   (import "zinkc" "emit_abi" (func (;2;) (type 1)))
   (import "env" "memory" (memory (;0;) 17))
   (func (;3;) (type 2) (param i32) (result i32)
-    i32.const 0
     local.get 0
+    i32.const 0
     call 0
     i32.const 0
     call 1)
@@ -19,8 +19,8 @@
     i32.const 34
     call 2)
   (func (;5;) (type 3) (param i32)
-    i32.const 0
     local.get 0
+    i32.const 0
     call 0)
   (func (;6;) (type 0)
     i32.const 1048610

diff --git a/compiler/wat/storage/load.wat b/compiler/wat/storage/load.wat
@@ -4,8 +4,8 @@
   (import "evm" "sload" (func (;0;) (type 0)))
   (import "evm" "sstore" (func (;1;) (type 1)))
   (func (type 0) (param i32) (result i32)
-        i32.const 0
         local.get 0
+        i32.const 0
         call 1
         i32.const 0
         call 0))
diff --git a/compiler/wat/storage/store.wat b/compiler/wat/storage/store.wat
@@ -2,7 +2,7 @@
   (type (;0;) (func (param i32 i32)))
   (import "evm" "sstore" (func (;0;) (type 0)))
   (func (param i32)
-        i32.const 0
         local.get 0
+        i32.const 0
         call 0)
  )