refactor(codegen): introduce WASM environment

codegen/src/codegen/constructor.rs

@@ -1,9 +1,7 @@
 //! Contract constructor.
 
-use crate::{
-    wasm::{self, ToLSBytes},
-    Buffer, Function, JumpTable, MacroAssembler, Result,
-};
+use crate::{wasm::ToLSBytes, Buffer, Function, JumpTable, MacroAssembler, Result};
+use wasmparser::FuncType;
 
 /// Contract constructor.
 ///
@@ -16,25 +14,22 @@ use crate::{
 ///
 /// TODO: introduce ABI for constructor
 pub struct Constructor {
-    /// Code buffer.
-    pub masm: MacroAssembler,
-
     /// Code generator.
+    pub masm: MacroAssembler,
+    /// Code buffer.
     pub init_code: Buffer,
-
     /// Runtime bytecode.
     pub runtime_bytecode: Buffer,
 }
 
 impl Constructor {
     /// Create a new constructor.
-    pub fn new(constructor: Option<wasm::Function<'_>>, runtime_bytecode: Buffer) -> Result<Self> {
+    pub fn new(constructor: Option<FuncType>, runtime_bytecode: Buffer) -> Result<Self> {
         let mut init_code = Buffer::new();
         if let Some(constructor) = constructor {
             let codegen = Function::new(
-                constructor.sig()?,
-                Default::default(),
                 Default::default(),
+                constructor,
                 // No `return` instruction in the generated code.
                 false,
             )?;

codegen/src/codegen/dispatcher.rs

@@ -2,7 +2,7 @@
 
 use crate::{
     codegen::code::ExtFunc,
-    wasm::{self, Data, Exports, Functions, Imports, ToLSBytes},
+    wasm::{self, Env, Functions, ToLSBytes},
     Error, JumpTable, MacroAssembler, Result,
 };
 use wasmparser::{FuncType, Operator};
@@ -12,14 +12,10 @@ use zabi::Abi;
 pub struct Dispatcher<'d> {
     /// Code buffer
     pub asm: MacroAssembler,
-    /// Module exports
-    pub exports: Exports,
+    /// WASM environment
+    pub env: Env,
     /// Module functions
     pub funcs: &'d Functions<'d>,
-    /// Module imports
-    pub imports: Imports,
-    /// Module data
-    pub data: Data,
     /// Jump table
     pub table: JumpTable,
     /// ABI for the current function
@@ -30,39 +26,19 @@ pub struct Dispatcher<'d> {
 
 impl<'d> Dispatcher<'d> {
     /// Create dispatcher with functions.
-    pub fn new(funcs: &'d Functions<'d>) -> Self {
+    pub fn new(env: Env, funcs: &'d Functions<'d>) -> Self {
         Self {
             asm: Default::default(),
-            exports: Default::default(),
+            env,
             funcs,
-            imports: Default::default(),
-            data: Default::default(),
             table: Default::default(),
             abi: Default::default(),
         }
     }
 
-    /// Set exports for the dispatcher.
-    pub fn exports(&mut self, exports: Exports) -> &mut Self {
-        self.exports = exports;
-        self
-    }
-
-    /// Set imports for the dispatcher.
-    pub fn imports(&mut self, imports: Imports) -> &mut Self {
-        self.imports = imports;
-        self
-    }
-
-    /// Set data for the dispatcher.
-    pub fn data(&mut self, data: Data) -> &mut Self {
-        self.data = data;
-        self
-    }
-
     /// Query exported function from selector.
     fn query_func(&self, name: &str) -> Result<u32> {
-        for (index, export) in self.exports.iter() {
+        for (index, export) in self.env.exports.iter() {
             if export == name {
                 return Ok(*index);
             }
@@ -90,11 +66,11 @@ impl<'d> Dispatcher<'d> {
             return Err(Error::InvalidSelector);
         };
 
-        if !self.imports.is_emit_abi(index) {
+        if !self.env.imports.is_emit_abi(index) {
             return Err(Error::FuncNotImported("emit_abi".into()));
         }
 
-        let abi = self.data.load(offset, length as usize)?;
+        let abi = self.env.data.load(offset, length as usize)?;
         Abi::from_hex(String::from_utf8_lossy(&abi)).map_err(Into::into)
     }
 

codegen/src/codegen/function.rs

@@ -6,50 +6,47 @@ use crate::{
     local::{LocalSlot, LocalSlotType, Locals},
     masm::MacroAssembler,
     validator::ValidateThenVisit,
-    wasm::{Data, Imports},
+    wasm::Env,
     Buffer, Error, Result,
 };
 use wasmparser::{FuncType, FuncValidator, LocalsReader, OperatorsReader, ValidatorResources};
 
 /// The code generation abstraction.
 pub struct Function {
     /// The backtrace.
-    pub(crate) backtrace: Backtrace,
+    pub backtrace: Backtrace,
     /// Control stack frames.
-    pub(crate) control: ControlStack,
-    /// Control stack frames.
-    pub(crate) dataset: Data,
-    /// The function environment.
-    pub(crate) env: FuncType,
+    pub control: ControlStack,
+    /// WASM environment.
+    pub env: Env,
     /// The defined locals for a function.
-    pub(crate) locals: Locals,
+    pub locals: Locals,
     /// The macro assembler.
-    pub(crate) masm: MacroAssembler,
+    pub masm: MacroAssembler,
     /// The jump table.
-    pub(crate) table: JumpTable,
-    /// The imported functions.
-    pub(crate) imports: Imports,
+    pub table: JumpTable,
+    /// The function type.
+    pub ty: FuncType,
     /// If this function is the main function.
-    pub(crate) is_main: bool,
+    pub is_main: bool,
 }
 
 impl Function {
     /// Create a new code generator.
-    pub fn new(env: FuncType, dataset: Data, imports: Imports, is_main: bool) -> Result<Self> {
+    pub fn new(env: Env, ty: FuncType, is_main: bool) -> Result<Self> {
         let mut params_count = 0;
         if !is_main {
-            params_count = env.params().len() as u8;
+            params_count = ty.params().len() as u8;
         }
 
         let mut codegen = Self {
             backtrace: Backtrace::default(),
             control: ControlStack::default(),
-            dataset,
             env,
+            ty,
             locals: Default::default(),
             masm: Default::default(),
             table: Default::default(),
-            imports,
             is_main,
         };
 
@@ -81,7 +78,7 @@ impl Function {
         let mut sp = if self.is_main { 0 } else { 1 };
 
         // Define locals in function parameters.
-        for param in self.env.params() {
+        for param in self.ty.params() {
             self.locals
                 .push(LocalSlot::new(*param, LocalSlotType::Parameter, sp));
             sp += 1;
@@ -128,7 +125,7 @@ impl Function {
     /// Finish code generation.
     pub fn finish(self, jump_table: &mut JumpTable, pc: u16) -> Result<Buffer> {
         let sp = self.masm.sp();
-        if !self.is_main && self.masm.sp() != self.env.results().len() as u8 {
+        if !self.is_main && self.masm.sp() != self.ty.results().len() as u8 {
             return Err(Error::StackNotBalanced(sp));
         }
 

codegen/src/visitor/call.rs

@@ -16,7 +16,7 @@ impl Function {
 
     /// The call instruction calls a function specified by its index.
     pub fn _call(&mut self, index: u32) -> Result<()> {
-        if self.imports.len() as u32 > index {
+        if self.env.imports.len() as u32 > index {
             self.call_imported(index)
         } else {
             self.call_internal(index)
@@ -55,6 +55,7 @@ impl Function {
         //
         // register the imported function index to the jump table.
         let func = *self
+            .env
             .imports
             .get(&index)
             .ok_or(Error::ImportedFuncNotFound(index))?;

codegen/src/visitor/handlers.rs

@@ -5,15 +5,15 @@ use crate::{ControlStackFrame, ControlStackFrameType, Function, Result};
 impl Function {
     /// Handle the end of the function.
     pub(crate) fn handle_return(&mut self) -> Result<()> {
-        let results = self.env.results();
+        let results = self.ty.results();
         tracing::trace!("handle return, results: {results:?}");
 
         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();
+        let results = self.ty.results();
         tracing::trace!("handle call return: {:?}", results);
 
         self.masm.call_return(results)

codegen/src/visitor/local.rs

@@ -6,7 +6,7 @@ impl Function {
     /// This instruction gets the value of a variable.
     pub fn _local_get(&mut self, local_index: u32) -> Result<()> {
         let local_index = local_index as usize;
-        if self.is_main && local_index < self.env.params().len() {
+        if self.is_main && local_index < self.ty.params().len() {
             self._local_get_calldata(local_index)
         } else {
             self._local_get_var(local_index)

codegen/src/visitor/log.rs

@@ -55,7 +55,7 @@ impl Function {
         for topic in (1..=count).rev() {
             let (offset, size) = self.log_data()?;
             let size = size as usize;
-            let data = self.dataset.load(offset, size)?;
+            let data = self.env.data.load(offset, size)?;
 
             tracing::debug!("log{count} topic{topic}: {:?}", data);
             topics.push(data);
@@ -64,7 +64,7 @@ impl Function {
         let name = {
             let (offset, size) = self.log_data()?;
             let size = size as usize;
-            let data = self.dataset.load(offset, size)?;
+            let data = self.env.data.load(offset, size)?;
 
             tracing::debug!("log1 name: {:?}", data);
             data

codegen/src/wasm/func.rs

@@ -34,24 +34,10 @@ impl Function<'_> {
     }
 }
 
-/// Functions with indexes.
+/// WASM Functions by indexes.
 #[derive(Default)]
 pub struct Functions<'f>(BTreeMap<u32, Function<'f>>);
 
-impl<'f> Deref for Functions<'f> {
-    type Target = BTreeMap<u32, Function<'f>>;
-
-    fn deref(&self) -> &Self::Target {
-        &self.0
-    }
-}
-
-impl<'f> DerefMut for Functions<'f> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.0
-    }
-}
-
 impl<'f> Functions<'f> {
     /// Add function to the list.
     pub fn add(
@@ -69,10 +55,10 @@ impl<'f> Functions<'f> {
     }
 
     /// Remove constructor function
-    pub fn remove_constructor(&mut self, exports: &Exports) -> Option<Function<'f>> {
+    pub fn remove_constructor(&mut self, exports: &Exports) -> Option<FuncType> {
         for (index, export) in exports.iter() {
             if export.as_str() == "constructor" {
-                return self.remove(index);
+                return self.remove(index)?.sig().ok();
             }
         }
 
@@ -97,3 +83,17 @@ impl<'f> Functions<'f> {
         self.0.into_values().collect()
     }
 }
+
+impl<'f> Deref for Functions<'f> {
+    type Target = BTreeMap<u32, Function<'f>>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<'f> DerefMut for Functions<'f> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}