chore: refactor FunctionInstance

Signed-off-by: Henry <mail@henrygressmann.de>

6 files changed, 62 insertions, 66 deletions

diff --git a/crates/tinywasm/src/func.rs b/crates/tinywasm/src/func.rs
index 18bcff9..aa7c674 100644
--- a/crates/tinywasm/src/func.rs
+++ b/crates/tinywasm/src/func.rs
@@ -1,6 +1,6 @@
 use crate::interpreter::stack::CallFrame;
 use crate::reference::StoreItem;
-use crate::{Error, FunctionDef, InterpreterRuntime, Result, Store, unlikely};
+use crate::{Error, FunctionInstance, InterpreterRuntime, Result, Store, unlikely};
 use alloc::rc::Rc;
 use alloc::{boxed::Box, format, string::ToString, vec, vec::Vec};
 use tinywasm_types::{ExternRef, FuncRef, FuncType, ModuleInstanceAddr, ValType, WasmValue};
@@ -14,20 +14,19 @@ impl Function {
         self.item.validate_store(store)?;
         validate_call_params(&self.ty, params)?;
 
-        let func_inst = store.state.get_func(self.addr);
-        let wasm_func = match &func_inst.func {
-            FunctionDef::Host(host_func) => {
+        let wasm_func = match store.state.get_func(self.addr) {
+            FunctionInstance::Host(host_func) => {
                 return host_func.clone().call(FuncContext { store, module_addr: self.module_addr }, params);
             }
-            FunctionDef::Wasm(wasm_func) => wasm_func.clone(),
+            FunctionInstance::Wasm(wasm_func) => wasm_func,
         };
 
         // Reset stack, push args, allocate locals, create entry frame.
         store.stack.clear();
         store.stack.values.extend_from_wasmvalues(params)?;
-        let locals_base = store.stack.values.enter_locals(&wasm_func.params, &wasm_func.locals)?;
-        let stack_offset = wasm_func.locals;
-        let callframe = CallFrame::new(self.addr, func_inst.owner, locals_base, stack_offset);
+        let locals_base = store.stack.values.enter_locals(&wasm_func.func.params, &wasm_func.func.locals)?;
+        let stack_offset = wasm_func.func.locals;
+        let callframe = CallFrame::new(self.addr, wasm_func.owner, locals_base, stack_offset);
 
         // Execute until completion and then collect result values from the stack.
         InterpreterRuntime::exec(store, callframe)?;
@@ -48,21 +47,17 @@ impl Function {
         self.item.validate_store(store)?;
         validate_call_params(&self.ty, params)?;
 
-        let func_inst = store.state.get_func(self.addr);
-        let func_inst_owner = func_inst.owner;
-        let func = func_inst.func.clone();
-
-        match func {
-            FunctionDef::Host(host_func) => {
-                let result = host_func.call(FuncContext { store, module_addr: self.module_addr }, params)?;
+        match store.state.get_func(self.addr) {
+            FunctionInstance::Host(host_func) => {
+                let result = host_func.clone().call(FuncContext { store, module_addr: self.module_addr }, params)?;
                 Ok(FuncExecution { store, state: FuncExecutionState::Completed { result: Some(result) } })
             }
-            FunctionDef::Wasm(wasm_func) => {
+            FunctionInstance::Wasm(wasm_func) => {
                 store.stack.clear();
                 store.stack.values.extend_from_wasmvalues(params)?;
-                let locals_base = store.stack.values.enter_locals(&wasm_func.params, &wasm_func.locals)?;
-                let stack_offset = wasm_func.locals;
-                let callframe = CallFrame::new(self.addr, func_inst_owner, locals_base, stack_offset);
+                let locals_base = store.stack.values.enter_locals(&wasm_func.func.params, &wasm_func.func.locals)?;
+                let stack_offset = wasm_func.func.locals;
+                let callframe = CallFrame::new(self.addr, wasm_func.owner, locals_base, stack_offset);
 
                 Ok(FuncExecution {
                     store,
@@ -151,7 +146,7 @@ impl HostFunction {
             Ok(result)
         };
 
-        let addr = store.add_func(FunctionDef::Host(Rc::new(Self { func: Box::new(inner_func), ty: ty.clone() })), 0);
+        let addr = store.add_func(FunctionInstance::Host(Rc::new(Self { func: Box::new(inner_func), ty: ty.clone() })));
         Function { item: crate::StoreItem::new(store.id(), addr), module_addr: 0, addr, ty: ty.clone() }
     }
 
@@ -169,7 +164,7 @@ impl HostFunction {
 
         let results = R::val_types();
         let ty = tinywasm_types::FuncType { params: P::val_types(), results };
-        let addr = store.add_func(FunctionDef::Host(Rc::new(Self { func: Box::new(inner_func), ty: ty.clone() })), 0);
+        let addr = store.add_func(FunctionInstance::Host(Rc::new(Self { func: Box::new(inner_func), ty: ty.clone() })));
         Function { item: crate::StoreItem::new(store.id(), addr), module_addr: 0, addr, ty }
     }
 }
@@ -317,7 +312,7 @@ impl<'store> FuncExecution<'store> {
 
         match InterpreterRuntime::exec_with_fuel(self.store, exec_state.callframe, fuel)? {
             crate::interpreter::ExecState::Completed => {
-                let result_ty = self.store.state.get_func(*root_func_addr).func.ty().clone();
+                let result_ty = self.store.state.get_func(*root_func_addr).ty().clone();
                 let result = collect_call_results(self.store, &result_ty)?;
                 self.state = FuncExecutionState::Completed { result: None };
                 Ok(ExecProgress::Completed(result))
@@ -354,7 +349,7 @@ impl<'store> FuncExecution<'store> {
 
         match InterpreterRuntime::exec_with_time_budget(self.store, exec_state.callframe, time_budget)? {
             crate::interpreter::ExecState::Completed => {
-                let result_ty = self.store.state.get_func(*root_func_addr).func.ty().clone();
+                let result_ty = self.store.state.get_func(*root_func_addr).ty().clone();
                 let result = collect_call_results(self.store, &result_ty)?;
                 self.state = FuncExecutionState::Completed { result: None };
                 Ok(ExecProgress::Completed(result))
diff --git a/crates/tinywasm/src/imports.rs b/crates/tinywasm/src/imports.rs
index b8dee06..abd9b70 100644
--- a/crates/tinywasm/src/imports.rs
+++ b/crates/tinywasm/src/imports.rs
@@ -339,7 +339,7 @@ impl Imports {
                             .get(*ty as usize)
                             .ok_or_else(|| LinkingError::incompatible_import_type(import))?;
 
-                        Self::compare_types(import, func.func.ty(), import_func_type)?;
+                        Self::compare_types(import, func.ty(), import_func_type)?;
                         imports.funcs.push(func_addr);
                     }
                     _ => return Err(LinkingError::incompatible_import_type(import).into()),
diff --git a/crates/tinywasm/src/instance.rs b/crates/tinywasm/src/instance.rs
index 12148db..2e5b4b6 100644
--- a/crates/tinywasm/src/instance.rs
+++ b/crates/tinywasm/src/instance.rs
@@ -273,7 +273,7 @@ impl ModuleInstance {
             return Err(Error::Other(format!("Export is not a function: {name}")));
         };
 
-        let ty = store.state.get_func(func_addr).func.ty();
+        let ty = store.state.get_func(func_addr).ty();
         Ok(Function {
             item: crate::StoreItem::new(self.0.store_id, func_addr),
             addr: func_addr,
@@ -294,7 +294,7 @@ impl ModuleInstance {
         self.validate_store(store)?;
         let func_addr = Self::index_addr(&self.0.func_addrs, func_index, "function")?;
 
-        let ty = store.state.get_func(func_addr).func.ty();
+        let ty = store.state.get_func(func_addr).ty();
         Ok(Function {
             item: crate::StoreItem::new(self.0.store_id, func_addr),
             addr: func_addr,
@@ -437,7 +437,7 @@ impl ModuleInstance {
         };
 
         let func_addr = self.0.resolve_func_addr(func_index);
-        let ty = store.state.get_func(func_addr).func.ty();
+        let ty = store.state.get_func(func_addr).ty();
         Ok(Some(Function {
             item: crate::StoreItem::new(self.0.store_id, func_addr),
             module_addr: self.id(),
diff --git a/crates/tinywasm/src/interpreter/executor.rs b/crates/tinywasm/src/interpreter/executor.rs
index d4dafd6..e76220f 100644
--- a/crates/tinywasm/src/interpreter/executor.rs
+++ b/crates/tinywasm/src/interpreter/executor.rs
@@ -710,21 +710,20 @@ impl<'store, const BUDGETED: bool> Executor<'store, BUDGETED> {
 
     fn exec_call<const IS_RETURN_CALL: bool>(
         &mut self,
-        wasm_func: Rc<WasmFunction>,
+        wasm_func: WasmFunctionInstance,
         func_addr: FuncAddr,
-        owner: ModuleInstanceAddr,
     ) -> Result<()> {
-        if !Rc::ptr_eq(&self.func, &wasm_func) {
-            self.func = wasm_func.clone();
+        if !Rc::ptr_eq(&self.func, &wasm_func.func) {
+            self.func = wasm_func.func.clone();
         }
 
         if IS_RETURN_CALL {
-            self.store.stack.values.truncate_keep_counts(self.cf.locals_base, wasm_func.params);
+            self.store.stack.values.truncate_keep_counts(self.cf.locals_base, wasm_func.func.params);
         }
 
-        let res = self.store.stack.values.enter_locals(&wasm_func.params, &wasm_func.locals);
+        let res = self.store.stack.values.enter_locals(&wasm_func.func.params, &wasm_func.func.locals);
         let locals_base = res.map_err(|err| if IS_RETURN_CALL { err } else { Error::Trap(Trap::CallStackOverflow) })?;
-        let new_call_frame = CallFrame::new(func_addr, owner, locals_base, wasm_func.locals);
+        let new_call_frame = CallFrame::new(func_addr, wasm_func.owner, locals_base, wasm_func.func.locals);
 
         if !IS_RETURN_CALL {
             self.cf.incr_instr_ptr(); // skip the call instruction
@@ -748,12 +747,9 @@ impl<'store, const BUDGETED: bool> Executor<'store, BUDGETED> {
     fn exec_call_direct<const IS_RETURN_CALL: bool>(&mut self, v: u32) -> Result<()> {
         self.charge_call_fuel(FUEL_COST_CALL_TOTAL);
         let addr = self.module.resolve_func_addr(v);
-        let func_inst = self.store.state.get_func(addr);
-        match &func_inst.func {
-            crate::FunctionDef::Wasm(wasm_func) => {
-                self.exec_call::<IS_RETURN_CALL>(wasm_func.clone(), addr, func_inst.owner)
-            }
-            crate::FunctionDef::Host(host_func) => self.exec_call_host(host_func.clone()),
+        match self.store.state.get_func(addr) {
+            crate::FunctionInstance::Wasm(wasm_func) => self.exec_call::<IS_RETURN_CALL>(wasm_func.clone(), addr),
+            crate::FunctionInstance::Host(host_func) => self.exec_call_host(host_func.clone()),
         }
     }
 
@@ -790,22 +786,20 @@ impl<'store, const BUDGETED: bool> Executor<'store, BUDGETED> {
             table.addr().ok_or_else(|| Error::from(Trap::UninitializedElement { index: table_idx as usize }))?
         };
 
-        let func_inst = self.store.state.get_func(func_ref);
         let call_ty = self.module.func_ty(type_addr);
-
-        match &func_inst.func {
-            crate::FunctionDef::Wasm(wasm_func) => {
-                if wasm_func.ty != *call_ty {
+        match self.store.state.get_func(func_ref) {
+            crate::FunctionInstance::Wasm(wasm_func) => {
+                if wasm_func.ty() != call_ty {
                     return Err(Trap::IndirectCallTypeMismatch {
-                        actual: wasm_func.ty.clone(),
+                        actual: wasm_func.ty().clone(),
                         expected: call_ty.clone(),
                     }
                     .into());
                 }
 
-                self.exec_call::<IS_RETURN_CALL>(wasm_func.clone(), func_ref, func_inst.owner)
+                self.exec_call::<IS_RETURN_CALL>(wasm_func.clone(), func_ref)
             }
-            crate::FunctionDef::Host(host_func) => {
+            crate::FunctionInstance::Host(host_func) => {
                 if host_func.ty != *call_ty {
                     return Err(Trap::IndirectCallTypeMismatch {
                         actual: host_func.ty.clone(),
diff --git a/crates/tinywasm/src/store/function.rs b/crates/tinywasm/src/store/function.rs
index b8edf67..56479dc 100644
--- a/crates/tinywasm/src/store/function.rs
+++ b/crates/tinywasm/src/store/function.rs
@@ -8,33 +8,40 @@ use crate::func::HostFunction;
 /// See <https://webassembly.github.io/spec/core/exec/runtime.html#function-instances>
 #[derive(Clone)]
 #[cfg_attr(feature = "debug", derive(Debug))]
-pub(crate) struct FunctionInstance {
-    pub(crate) func: FunctionDef,
-    pub(crate) owner: ModuleInstanceAddr, // index into store.module_instances, none for host functions
-}
-
-/// The internal representation of a function
-#[derive(Clone)]
-#[cfg_attr(feature = "debug", derive(Debug))]
-pub(crate) enum FunctionDef {
+pub(crate) enum FunctionInstance {
     /// A host function
     Host(Rc<HostFunction>),
 
     /// A pointer to a WebAssembly function
-    Wasm(Rc<WasmFunction>),
+    Wasm(WasmFunctionInstance),
 }
 
-impl FunctionDef {
+impl FunctionInstance {
+    #[inline]
     pub(crate) fn ty(&self) -> &FuncType {
         match self {
             Self::Host(f) => &f.ty,
-            Self::Wasm(f) => &f.ty,
+            Self::Wasm(f) => f.ty(),
         }
     }
 }
 
 impl FunctionInstance {
     pub(crate) fn new_wasm(func: WasmFunction, owner: ModuleInstanceAddr) -> Self {
-        Self { func: FunctionDef::Wasm(Rc::new(func)), owner }
+        Self::Wasm(WasmFunctionInstance { func: Rc::new(func), owner })
+    }
+}
+
+#[derive(Clone)]
+#[cfg_attr(feature = "debug", derive(Debug))]
+pub(crate) struct WasmFunctionInstance {
+    pub(crate) func: Rc<WasmFunction>,
+    pub(crate) owner: ModuleInstanceAddr,
+}
+
+impl WasmFunctionInstance {
+    #[inline]
+    pub(crate) fn ty(&self) -> &FuncType {
+        &self.func.ty
     }
 }
diff --git a/crates/tinywasm/src/store/mod.rs b/crates/tinywasm/src/store/mod.rs
index f7bd0e9..4ba6c97 100644
--- a/crates/tinywasm/src/store/mod.rs
+++ b/crates/tinywasm/src/store/mod.rs
@@ -117,9 +117,9 @@ impl State {
     /// Get a wasm function at the actual index in the store, panicking if it's a host function (which should be guaranteed by the validator)
     pub(crate) fn get_wasm_func(&self, addr: FuncAddr) -> &Rc<WasmFunction> {
         match self.funcs.get(addr as usize) {
-            Some(func) => match &func.func {
-                FunctionDef::Wasm(wasm_func) => wasm_func,
-                FunctionDef::Host(_) => unreachable!(
+            Some(func) => match func {
+                FunctionInstance::Wasm(wasm_func) => &wasm_func.func,
+                FunctionInstance::Host(_) => unreachable!(
                     "expected a wasm function at address {addr}, but found a host function. This should be unreachable"
                 ),
             },
@@ -423,8 +423,8 @@ impl Store {
         Ok((data_addrs.into_boxed_slice(), None))
     }
 
-    pub(crate) fn add_func(&mut self, func: FunctionDef, idx: ModuleInstanceAddr) -> FuncAddr {
-        self.state.funcs.push(FunctionInstance { func, owner: idx });
+    pub(crate) fn add_func(&mut self, func: FunctionInstance) -> FuncAddr {
+        self.state.funcs.push(func);
         self.state.funcs.len() as FuncAddr - 1
     }