From dc08a9532bd10447024ff6dd28bc2ba9c3cda3c6 Mon Sep 17 00:00:00 2001
From: clearloop <tianyi.gc@gmail.com>
Date: Thu, 28 Dec 2023 17:00:18 +0800
Subject: [PATCH] feat(codegen): remove lifetime for dispatcher

---
 codegen/src/codegen/constructor.rs | 24 +++++------
 codegen/src/codegen/dispatcher.rs  | 69 ++++++++++++++++--------------
 compiler/src/compiler.rs           |  6 +--
 3 files changed, 52 insertions(+), 47 deletions(-)

diff --git a/codegen/src/codegen/constructor.rs b/codegen/src/codegen/constructor.rs
index 0721d356c..92a5dd5b1 100644
--- a/codegen/src/codegen/constructor.rs
+++ b/codegen/src/codegen/constructor.rs
@@ -46,18 +46,6 @@ impl Constructor {
         })
     }
 
-    /// Returns the length of instructions.
-    fn return_instr_length(init_code_length: usize, runtime_bytecode_length: usize) -> usize {
-        let mut expected_length =
-            runtime_bytecode_length.to_ls_bytes().len() + init_code_length.to_ls_bytes().len() + 3;
-
-        if init_code_length < 0xff && init_code_length + expected_length > 0xff {
-            expected_length += 1;
-        }
-
-        expected_length
-    }
-
     /// Concat the constructor code.
     ///
     /// Here we override the memory totally with
@@ -115,4 +103,16 @@ impl Constructor {
 
         Ok(self.masm.buffer().into())
     }
+
+    /// Returns the length of instructions.
+    fn return_instr_length(init_code_length: usize, runtime_bytecode_length: usize) -> usize {
+        let mut expected_length =
+            runtime_bytecode_length.to_ls_bytes().len() + init_code_length.to_ls_bytes().len() + 3;
+
+        if init_code_length < 0xff && init_code_length + expected_length > 0xff {
+            expected_length += 1;
+        }
+
+        expected_length
+    }
 }
diff --git a/codegen/src/codegen/dispatcher.rs b/codegen/src/codegen/dispatcher.rs
index f3ad8d239..3f4c5cce6 100644
--- a/codegen/src/codegen/dispatcher.rs
+++ b/codegen/src/codegen/dispatcher.rs
@@ -1,5 +1,7 @@
 //! Code generator for EVM dispatcher.
 
+use std::collections::BTreeMap;
+
 use crate::{
     codegen::code::ExtFunc,
     wasm::{self, Env, Functions, ToLSBytes},
@@ -9,31 +11,55 @@ use wasmparser::{FuncType, Operator};
 use zabi::Abi;
 
 /// Code generator for EVM dispatcher.
-pub struct Dispatcher<'d> {
+pub struct Dispatcher {
+    /// ABI for the current function
+    pub abi: Vec<Abi>,
     /// Code buffer
     pub asm: MacroAssembler,
     /// WASM environment
     pub env: Env,
     /// Module functions
-    pub funcs: &'d Functions<'d>,
+    pub funcs: BTreeMap<u32, FuncType>,
     /// Jump table
     pub table: JumpTable,
-    /// ABI for the current function
-    ///
-    /// TODO: refactor this. (#206)
-    pub abi: Vec<Abi>,
 }
 
-impl<'d> Dispatcher<'d> {
+impl Dispatcher {
     /// Create dispatcher with functions.
-    pub fn new(env: Env, funcs: &'d Functions<'d>) -> Self {
-        Self {
+    pub fn new(env: Env, funcs: &Functions<'_>) -> Result<Self> {
+        let funcs = funcs
+            .values()
+            .map(|func| Ok((func.index(), func.sig()?)))
+            .collect::<Result<_>>()?;
+
+        Ok(Self {
+            abi: Default::default(),
             asm: Default::default(),
             env,
             funcs,
             table: Default::default(),
-            abi: Default::default(),
+        })
+    }
+
+    /// Emit compiled code to the given buffer.
+    pub fn finish(&mut self, selectors: Functions<'_>, table: &mut JumpTable) -> Result<Vec<u8>> {
+        if selectors.is_empty() {
+            return Err(Error::SelectorNotFound);
+        }
+
+        self.asm._push0()?;
+        self.asm._calldataload()?;
+        self.asm.push(&[0xe0])?;
+        self.asm._shr()?;
+
+        let mut len = selectors.len();
+        for (_, func) in selectors.iter() {
+            self.emit_selector(func, len == 1)?;
+            len -= 1;
         }
+
+        table.merge(self.table.clone(), 0)?;
+        Ok(self.asm.buffer().into())
     }
 
     /// Query exported function from selector.
@@ -178,7 +204,7 @@ impl<'d> Dispatcher<'d> {
             .funcs
             .get(&func)
             .ok_or(Error::FuncNotFound(func))?
-            .sig()?;
+            .clone();
 
         // TODO: optimize this on parameter length (#165)
         {
@@ -220,25 +246,4 @@ impl<'d> Dispatcher<'d> {
 
         Ok(())
     }
-
-    /// Emit compiled code to the given buffer.
-    pub fn finish(&mut self, selectors: Functions<'_>, table: &mut JumpTable) -> Result<Vec<u8>> {
-        if selectors.is_empty() {
-            return Err(Error::SelectorNotFound);
-        }
-
-        self.asm._push0()?;
-        self.asm._calldataload()?;
-        self.asm.push(&[0xe0])?;
-        self.asm._shr()?;
-
-        let mut len = selectors.len();
-        for (_, func) in selectors.iter() {
-            self.emit_selector(func, len == 1)?;
-            len -= 1;
-        }
-
-        table.merge(self.table.clone(), 0)?;
-        Ok(self.asm.buffer().into())
-    }
 }
diff --git a/compiler/src/compiler.rs b/compiler/src/compiler.rs
index b014cfb84..0bc1db1c0 100644
--- a/compiler/src/compiler.rs
+++ b/compiler/src/compiler.rs
@@ -14,10 +14,10 @@ pub struct Compiler {
     pub(crate) abi: Vec<Abi>,
     /// EVM bytecode buffer.
     pub(crate) buffer: Buffer,
-    /// Global jump table.
-    table: JumpTable,
     /// Compiler configuration.
     pub config: Config,
+    /// Global jump table.
+    table: JumpTable,
 }
 
 impl Compiler {
@@ -58,7 +58,7 @@ impl Compiler {
             return Ok(());
         }
 
-        let mut dispatcher = Dispatcher::new(parser.to_env(), &parser.funcs);
+        let mut dispatcher = Dispatcher::new(parser.to_env(), &parser.funcs)?;
         let buffer = dispatcher.finish(selectors, &mut self.table)?;
         self.buffer.extend_from_slice(&buffer);