chore: simplify interpreter::exec

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

9 files changed, 128 insertions, 150 deletions

diff --git a/benches/fibonacci.rs b/benches/fibonacci.rs
index b0a4834..7ccde02 100644
--- a/benches/fibonacci.rs
+++ b/benches/fibonacci.rs
@@ -30,7 +30,7 @@ fn criterion_benchmark(c: &mut Criterion) {
 
     let mut group = c.benchmark_group("fibonacci");
     group.bench_function("tinywasm", |b| b.iter(|| run_tinywasm(module.clone())));
-    // group.bench_function("wasmi", |b| b.iter(|| run_wasmi()));
+    group.bench_function("wasmi", |b| b.iter(|| run_wasmi()));
 }
 
 criterion_group!(
diff --git a/crates/tinywasm/src/func.rs b/crates/tinywasm/src/func.rs
index 1c5129b..a0c1212 100644
--- a/crates/tinywasm/src/func.rs
+++ b/crates/tinywasm/src/func.rs
@@ -1,17 +1,17 @@
-use crate::{log, runtime::RawWasmValue};
+use crate::{log, runtime::RawWasmValue, unlikely, Function};
 use alloc::{boxed::Box, format, string::String, string::ToString, vec, vec::Vec};
-use tinywasm_types::{FuncAddr, FuncType, ValType, WasmValue};
+use tinywasm_types::{FuncType, ModuleInstanceAddr, ValType, WasmValue};
 
 use crate::{
     runtime::{CallFrame, Stack},
-    Error, FuncContext, ModuleInstance, Result, Store,
+    Error, FuncContext, Result, Store,
 };
 
 #[derive(Debug)]
 /// A function handle
 pub struct FuncHandle {
-    pub(crate) module: ModuleInstance,
-    pub(crate) addr: FuncAddr,
+    pub(crate) module_addr: ModuleInstanceAddr,
+    pub(crate) addr: u32,
     pub(crate) ty: FuncType,
 
     /// The name of the function, if it has one
@@ -22,18 +22,13 @@ impl FuncHandle {
     /// Call a function (Invocation)
     ///
     /// See <https://webassembly.github.io/spec/core/exec/modules.html#invocation>
+    #[inline]
     pub fn call(&self, store: &mut Store, params: &[WasmValue]) -> Result<Vec<WasmValue>> {
-        let mut stack = Stack::default();
-
-        // 1. Assert: funcs[func_addr] exists
-        // 2. let func_inst be the functiuon instance funcs[func_addr]
-        let func_inst = store.get_func(self.addr as usize)?.clone();
-
         // 3. Let func_ty be the function type
         let func_ty = &self.ty;
 
         // 4. If the length of the provided argument values is different from the number of expected arguments, then fail
-        if func_ty.params.len() != params.len() {
+        if unlikely(func_ty.params.len() != params.len()) {
             log::info!("func_ty.params: {:?}", func_ty.params);
             return Err(Error::Other(format!(
                 "param count mismatch: expected {}, got {}",
@@ -43,31 +38,34 @@ impl FuncHandle {
         }
 
         // 5. For each value type and the corresponding value, check if types match
-        for (i, (ty, param)) in func_ty.params.iter().zip(params).enumerate() {
+        if !unlikely(func_ty.params.iter().zip(params).enumerate().all(|(i, (ty, param))| {
             if ty != &param.val_type() {
-                return Err(Error::Other(format!(
-                    "param type mismatch at index {}: expected {:?}, got {:?}",
-                    i, ty, param
-                )));
+                log::error!("param type mismatch at index {}: expected {:?}, got {:?}", i, ty, param);
+                false
+            } else {
+                true
             }
+        })) {
+            return Err(Error::Other("Type mismatch".into()));
         }
 
-        let locals = match &func_inst.func {
-            crate::Function::Host(h) => {
-                let func = h.func.clone();
-                let ctx = FuncContext { store, module: &self.module };
+        let func_inst = store.get_func(self.addr as usize)?;
+        let wasm_func = match &func_inst.func {
+            Function::Host(host_func) => {
+                let func = &host_func.clone().func;
+                let ctx = FuncContext { store, module_addr: self.module_addr };
                 return (func)(ctx, params);
             }
-            crate::Function::Wasm(ref f) => f.locals.to_vec(),
+            Function::Wasm(wasm_func) => wasm_func,
         };
 
         // 6. Let f be the dummy frame
-        log::debug!("locals: {:?}", locals);
-        let call_frame = CallFrame::new(func_inst, params.iter().map(|v| RawWasmValue::from(*v)), locals);
+        let call_frame =
+            CallFrame::new(wasm_func.clone(), func_inst.owner, params.iter().map(|v| RawWasmValue::from(*v)));
 
         // 7. Push the frame f to the call stack
         // & 8. Push the values to the stack (Not needed since the call frame owns the values)
-        stack.call_stack.push(call_frame)?;
+        let mut stack = Stack::new(call_frame);
 
         // 9. Invoke the function instance
         let runtime = store.runtime();
@@ -125,6 +123,7 @@ macro_rules! impl_into_wasm_value_tuple {
             $($T: Into<WasmValue>),*
         {
             #[allow(non_snake_case)]
+            #[inline]
             fn into_wasm_value_tuple(self) -> Vec<WasmValue> {
                 let ($($T,)*) = self;
                 vec![$($T.into(),)*]
@@ -136,6 +135,7 @@ macro_rules! impl_into_wasm_value_tuple {
 macro_rules! impl_into_wasm_value_tuple_single {
     ($T:ident) => {
         impl IntoWasmValueTuple for $T {
+            #[inline]
             fn into_wasm_value_tuple(self) -> Vec<WasmValue> {
                 vec![self.into()]
             }
@@ -164,6 +164,7 @@ macro_rules! impl_from_wasm_value_tuple {
         where
             $($T: TryFrom<WasmValue, Error = ()>),*
         {
+            #[inline]
             fn from_wasm_value_tuple(values: &[WasmValue]) -> Result<Self> {
                 #[allow(unused_variables, unused_mut)]
                 let mut iter = values.iter();
@@ -186,6 +187,7 @@ macro_rules! impl_from_wasm_value_tuple {
 macro_rules! impl_from_wasm_value_tuple_single {
     ($T:ident) => {
         impl FromWasmValueTuple for $T {
+            #[inline]
             fn from_wasm_value_tuple(values: &[WasmValue]) -> Result<Self> {
                 #[allow(unused_variables, unused_mut)]
                 let mut iter = values.iter();
@@ -254,6 +256,7 @@ macro_rules! impl_val_types_from_tuple {
         where
             $($t: ToValType,)+
         {
+            #[inline]
             fn val_types() -> Box<[ValType]> {
                 Box::new([$($t::to_val_type(),)+])
             }
@@ -262,6 +265,7 @@ macro_rules! impl_val_types_from_tuple {
 }
 
 impl ValTypesFromTuple for () {
+    #[inline]
     fn val_types() -> Box<[ValType]> {
         Box::new([])
     }
@@ -271,6 +275,7 @@ impl<T1> ValTypesFromTuple for T1
 where
     T1: ToValType,
 {
+    #[inline]
     fn val_types() -> Box<[ValType]> {
         Box::new([T1::to_val_type()])
     }
diff --git a/crates/tinywasm/src/imports.rs b/crates/tinywasm/src/imports.rs
index 32140a7..0b16970 100644
--- a/crates/tinywasm/src/imports.rs
+++ b/crates/tinywasm/src/imports.rs
@@ -7,6 +7,7 @@ use crate::{
     log, LinkingError, Result,
 };
 use alloc::{
+    boxed::Box,
     collections::BTreeMap,
     rc::Rc,
     string::{String, ToString},
@@ -18,10 +19,10 @@ use tinywasm_types::*;
 #[derive(Debug, Clone)]
 pub enum Function {
     /// A host function
-    Host(HostFunction),
+    Host(Rc<HostFunction>),
 
     /// A function defined in WebAssembly
-    Wasm(WasmFunction),
+    Wasm(Rc<WasmFunction>),
 }
 
 impl Function {
@@ -34,7 +35,6 @@ impl Function {
 }
 
 /// A host function
-#[derive(Clone)]
 pub struct HostFunction {
     pub(crate) ty: tinywasm_types::FuncType,
     pub(crate) func: HostFuncInner,
@@ -52,13 +52,13 @@ impl HostFunction {
     }
 }
 
-pub(crate) type HostFuncInner = Rc<dyn Fn(FuncContext<'_>, &[WasmValue]) -> Result<Vec<WasmValue>>>;
+pub(crate) type HostFuncInner = Box<dyn Fn(FuncContext<'_>, &[WasmValue]) -> Result<Vec<WasmValue>>>;
 
 /// The context of a host-function call
 #[derive(Debug)]
 pub struct FuncContext<'a> {
     pub(crate) store: &'a mut crate::Store,
-    pub(crate) module: &'a crate::ModuleInstance,
+    pub(crate) module_addr: ModuleInstanceAddr,
 }
 
 impl FuncContext<'_> {
@@ -73,13 +73,13 @@ impl FuncContext<'_> {
     }
 
     /// Get a reference to the module instance
-    pub fn module(&self) -> &crate::ModuleInstance {
-        self.module
+    pub fn module(&self) -> crate::ModuleInstance {
+        self.store.get_module_instance_raw(self.module_addr)
     }
 
     /// Get a reference to an exported memory
     pub fn memory(&mut self, name: &str) -> Result<crate::MemoryRef> {
-        self.module.exported_memory(self.store, name)
+        self.module().exported_memory(self.store, name)
     }
 }
 
@@ -140,7 +140,7 @@ impl Extern {
         ty: &tinywasm_types::FuncType,
         func: impl Fn(FuncContext<'_>, &[WasmValue]) -> Result<Vec<WasmValue>> + 'static,
     ) -> Self {
-        Self::Function(Function::Host(HostFunction { func: Rc::new(func), ty: ty.clone() }))
+        Self::Function(Function::Host(Rc::new(HostFunction { func: Box::new(func), ty: ty.clone() })))
     }
 
     /// Create a new typed function import
@@ -159,7 +159,7 @@ impl Extern {
 
         let ty = tinywasm_types::FuncType { params: P::val_types(), results: R::val_types() };
 
-        Self::Function(Function::Host(HostFunction { func: Rc::new(inner_func), ty }))
+        Self::Function(Function::Host(Rc::new(HostFunction { func: Box::new(inner_func), ty })))
     }
 
     pub(crate) fn kind(&self) -> ExternalKind {
diff --git a/crates/tinywasm/src/instance.rs b/crates/tinywasm/src/instance.rs
index b79f551..ecd6ce8 100644
--- a/crates/tinywasm/src/instance.rs
+++ b/crates/tinywasm/src/instance.rs
@@ -42,6 +42,10 @@ impl ModuleInstance {
         self.0 = other.0;
     }
 
+    pub(crate) fn swap_with(&mut self, other_addr: ModuleInstanceAddr, store: &mut Store) {
+        self.swap(store.get_module_instance_raw(other_addr))
+    }
+
     /// Get the module instance's address
     pub fn id(&self) -> ModuleInstanceAddr {
         self.0.idx
@@ -172,7 +176,7 @@ impl ModuleInstance {
         let func_inst = store.get_func(func_addr as usize)?;
         let ty = func_inst.func.ty();
 
-        Ok(FuncHandle { addr: func_addr, module: self.clone(), name: Some(name.to_string()), ty: ty.clone() })
+        Ok(FuncHandle { addr: func_addr, module_addr: self.id(), name: Some(name.to_string()), ty: ty.clone() })
     }
 
     /// Get a typed exported function by name
@@ -230,7 +234,7 @@ impl ModuleInstance {
         let func_inst = store.get_func(*func_addr as usize)?;
         let ty = func_inst.func.ty();
 
-        Ok(Some(FuncHandle { module: self.clone(), addr: *func_addr, ty: ty.clone(), name: None }))
+        Ok(Some(FuncHandle { module_addr: self.id(), addr: *func_addr, ty: ty.clone(), name: None }))
     }
 
     /// Invoke the start function of the module
diff --git a/crates/tinywasm/src/runtime/interpreter/mod.rs b/crates/tinywasm/src/runtime/interpreter/mod.rs
index a59b544..fdbe697 100644
--- a/crates/tinywasm/src/runtime/interpreter/mod.rs
+++ b/crates/tinywasm/src/runtime/interpreter/mod.rs
@@ -7,7 +7,7 @@ use crate::{
 use alloc::format;
 use alloc::{string::ToString, vec::Vec};
 use core::ops::{BitAnd, BitOr, BitXor, Neg};
-use tinywasm_types::{ElementKind, Instruction, ValType};
+use tinywasm_types::{ElementKind, ValType};
 
 #[cfg(not(feature = "std"))]
 mod no_std_floats;
@@ -28,51 +28,24 @@ impl InterpreterRuntime {
         // The current call frame, gets updated inside of exec_one
         let mut cf = stack.call_stack.pop()?;
 
-        let mut func_inst = cf.func_instance.clone();
-        let mut wasm_func = func_inst.assert_wasm()?;
-
         // The function to execute, gets updated from ExecResult::Call
-        let mut instrs = &wasm_func.instructions;
-        let mut instr_count = instrs.len();
-        let mut current_module = store.get_module_instance_raw(func_inst.owner);
+        let mut current_module = store.get_module_instance_raw(cf.func_instance.1);
 
         loop {
-            if unlikely(cf.instr_ptr >= instr_count) {
-                cold();
-                log::error!("instr_ptr out of bounds: {} >= {}", cf.instr_ptr, instr_count);
-                return Err(Error::Other(format!("instr_ptr out of bounds: {} >= {}", cf.instr_ptr, instr_count)));
-            }
-            let instr = &instrs[cf.instr_ptr];
-
-            match exec_one(&mut cf, instr, instrs, stack, store, &current_module)? {
+            match exec_one(&mut cf, stack, store, &current_module)? {
                 // Continue execution at the new top of the call stack
                 ExecResult::Call => {
                     cf = stack.call_stack.pop()?;
-                    func_inst = cf.func_instance.clone();
-                    wasm_func =
-                        func_inst.assert_wasm().map_err(|_| Error::Other("call expected wasm function".to_string()))?;
-                    instrs = &wasm_func.instructions;
-                    instr_count = instrs.len();
-
-                    if cf.func_instance.owner != current_module.id() {
-                        current_module.swap(
-                            store
-                                .get_module_instance(cf.func_instance.owner)
-                                .ok_or_else(|| Error::Other("call expected module instance".to_string()))?
-                                .clone(),
-                        );
+                    if cf.func_instance.1 != current_module.id() {
+                        current_module.swap_with(cf.func_instance.1, store);
                     }
-
-                    continue;
                 }
 
                 // return from the function
                 ExecResult::Return => return Ok(()),
 
                 // continue to the next instruction and increment the instruction pointer
-                ExecResult::Ok => {
-                    cf.instr_ptr += 1;
-                }
+                ExecResult::Ok => cf.instr_ptr += 1,
 
                 // trap the program
                 ExecResult::Trap(trap) => {
@@ -114,18 +87,17 @@ macro_rules! break_to {
 /// Run a single step of the interpreter
 /// A seperate function is used so later, we can more easily implement
 /// a step-by-step debugger (using generators once they're stable?)
-// TODO: perf: don't push then pop the call frame, just pass it via ExecResult::Call instead
 #[inline(always)] // this improves performance by more than 20% in some cases
-fn exec_one(
-    cf: &mut CallFrame,
-    instr: &Instruction,
-    instrs: &[Instruction],
-    stack: &mut Stack,
-    store: &mut Store,
-    module: &ModuleInstance,
-) -> Result<ExecResult> {
+fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &ModuleInstance) -> Result<ExecResult> {
+    let instrs = &cf.func_instance.0.instructions;
+    if unlikely(cf.instr_ptr >= instrs.len() || instrs.is_empty()) {
+        cold();
+        log::error!("instr_ptr out of bounds: {} >= {}", cf.instr_ptr, instrs.len());
+        return Err(Error::Other(format!("instr_ptr out of bounds: {} >= {}", cf.instr_ptr, instrs.len())));
+    }
+
     use tinywasm_types::Instruction::*;
-    match instr {
+    match &instrs[cf.instr_ptr] {
         Nop => { /* do nothing */ }
         Unreachable => {
             cold();
@@ -152,19 +124,19 @@ fn exec_one(
             let func_idx = module.resolve_func_addr(*v);
             let func_inst = store.get_func(func_idx as usize)?.clone();
 
-            let (locals, ty) = match &func_inst.func {
-                crate::Function::Wasm(ref f) => (f.locals.to_vec(), f.ty.clone()),
+            let wasm_func = match &func_inst.func {
+                crate::Function::Wasm(wasm_func) => wasm_func.clone(),
                 crate::Function::Host(host_func) => {
-                    let func = host_func.func.clone();
+                    let func = &host_func.func;
                     let params = stack.values.pop_params(&host_func.ty.params)?;
-                    let res = (func)(FuncContext { store, module }, &params)?;
+                    let res = (func)(FuncContext { store, module_addr: module.id() }, &params)?;
                     stack.values.extend_from_typed(&res);
                     return Ok(ExecResult::Ok);
                 }
             };
 
-            let params = stack.values.pop_n_rev(ty.params.len())?;
-            let call_frame = CallFrame::new(func_inst, params, locals);
+            let params = stack.values.pop_n_rev(wasm_func.ty.params.len())?;
+            let call_frame = CallFrame::new(wasm_func, func_inst.owner, params);
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
@@ -191,29 +163,37 @@ fn exec_one(
             };
 
             let func_inst = store.get_func(func_ref as usize)?.clone();
-            let func_ty = func_inst.func.ty();
             let call_ty = module.func_ty(*type_addr);
 
-            if unlikely(func_ty != call_ty) {
-                log::error!("indirect call type mismatch: {:?} != {:?}", func_ty, call_ty);
-                return Err(
-                    Trap::IndirectCallTypeMismatch { actual: func_ty.clone(), expected: call_ty.clone() }.into()
-                );
-            }
-
-            let locals = match &func_inst.func {
-                crate::Function::Wasm(ref f) => f.locals.to_vec(),
+            let wasm_func = match func_inst.func {
+                crate::Function::Wasm(ref f) => f.clone(),
                 crate::Function::Host(host_func) => {
-                    let func = host_func.func.clone();
-                    let params = stack.values.pop_params(&func_ty.params)?;
-                    let res = (func)(FuncContext { store, module }, &params)?;
+                    if unlikely(host_func.ty != *call_ty) {
+                        log::error!("indirect call type mismatch: {:?} != {:?}", host_func.ty, call_ty);
+                        return Err(Trap::IndirectCallTypeMismatch {
+                            actual: host_func.ty.clone(),
+                            expected: call_ty.clone(),
+                        }
+                        .into());
+                    }
+
+                    let host_func = host_func.clone();
+                    let params = stack.values.pop_params(&host_func.ty.params)?;
+                    let res = (host_func.func)(FuncContext { store, module_addr: module.id() }, &params)?;
                     stack.values.extend_from_typed(&res);
                     return Ok(ExecResult::Ok);
                 }
             };
 
-            let params = stack.values.pop_n_rev(func_ty.params.len())?;
-            let call_frame = CallFrame::new(func_inst, params, locals);
+            if unlikely(wasm_func.ty != *call_ty) {
+                log::error!("indirect call type mismatch: {:?} != {:?}", wasm_func.ty, call_ty);
+                return Err(
+                    Trap::IndirectCallTypeMismatch { actual: wasm_func.ty.clone(), expected: call_ty.clone() }.into()
+                );
+            }
+
+            let params = stack.values.pop_n_rev(wasm_func.ty.params.len())?;
+            let call_frame = CallFrame::new(wasm_func, func_inst.owner, params);
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
@@ -243,18 +223,16 @@ fn exec_one(
 
             // falsy value is on the top of the stack
             if let Some(else_offset) = else_offset {
-                cf.enter_label(
-                    LabelFrame::new(
-                        cf.instr_ptr + *else_offset,
-                        cf.instr_ptr + *end_offset,
-                        stack.values.len(), // - params,
-                        BlockType::Else,
-                        args,
-                        module,
-                    ),
-                    &mut stack.values,
+                let label = LabelFrame::new(
+                    cf.instr_ptr + *else_offset,
+                    cf.instr_ptr + *end_offset,
+                    stack.values.len(), // - params,
+                    BlockType::Else,
+                    args,
+                    module,
                 );
                 cf.instr_ptr += *else_offset;
+                cf.enter_label(label, &mut stack.values);
             } else {
                 cf.instr_ptr += *end_offset;
             }
diff --git a/crates/tinywasm/src/runtime/stack.rs b/crates/tinywasm/src/runtime/stack.rs
index 285d967..31c41f0 100644
--- a/crates/tinywasm/src/runtime/stack.rs
+++ b/crates/tinywasm/src/runtime/stack.rs
@@ -7,8 +7,14 @@ pub(crate) use blocks::{BlockType, LabelFrame};
 pub(crate) use call_stack::CallFrame;
 
 /// A WebAssembly Stack
-#[derive(Debug, Default)]
+#[derive(Debug)]
 pub struct Stack {
     pub(crate) values: ValueStack,
     pub(crate) call_stack: CallStack,
 }
+
+impl Stack {
+    pub(crate) fn new(call_frame: CallFrame) -> Self {
+        Self { values: ValueStack::default(), call_stack: CallStack::new(call_frame) }
+    }
+}
diff --git a/crates/tinywasm/src/runtime/stack/blocks.rs b/crates/tinywasm/src/runtime/stack/blocks.rs
index ff77cf8..f302e59 100644
--- a/crates/tinywasm/src/runtime/stack/blocks.rs
+++ b/crates/tinywasm/src/runtime/stack/blocks.rs
@@ -4,7 +4,7 @@ use tinywasm_types::BlockArgs;
 use crate::{unlikely, ModuleInstance};
 
 #[derive(Debug, Clone)]
-pub(crate) struct Labels(Vec<LabelFrame>);
+pub(crate) struct Labels(Vec<LabelFrame>); // TODO: maybe Box<[LabelFrame]> by analyzing the lable count when parsing the module?
 
 impl Labels {
     pub(crate) fn new() -> Self {
diff --git a/crates/tinywasm/src/runtime/stack/call_stack.rs b/crates/tinywasm/src/runtime/stack/call_stack.rs
index ebec142..9ba6ad2 100644
--- a/crates/tinywasm/src/runtime/stack/call_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/call_stack.rs
@@ -1,10 +1,10 @@
 use crate::unlikely;
 use crate::{
     runtime::{BlockType, RawWasmValue},
-    Error, FunctionInstance, Result, Trap,
+    Error, Result, Trap,
 };
 use alloc::{boxed::Box, rc::Rc, vec::Vec};
-use tinywasm_types::ValType;
+use tinywasm_types::{ModuleInstanceAddr, WasmFunction};
 
 use super::{blocks::Labels, LabelFrame};
 
@@ -17,13 +17,14 @@ pub(crate) struct CallStack {
     stack: Vec<CallFrame>,
 }
 
-impl Default for CallStack {
-    fn default() -> Self {
-        Self { stack: Vec::with_capacity(CALL_STACK_SIZE) }
+impl CallStack {
+    #[inline]
+    pub(crate) fn new(initial_frame: CallFrame) -> Self {
+        let mut stack = Self { stack: Vec::with_capacity(CALL_STACK_SIZE) };
+        stack.push(initial_frame).unwrap();
+        stack
     }
-}
 
-impl CallStack {
     #[inline]
     pub(crate) fn is_empty(&self) -> bool {
         self.stack.is_empty()
@@ -51,14 +52,13 @@ impl CallStack {
 pub(crate) struct CallFrame {
     pub(crate) instr_ptr: usize,
     // pub(crate) module: ModuleInstanceAddr,
-    pub(crate) func_instance: Rc<FunctionInstance>,
-
+    pub(crate) func_instance: (Rc<WasmFunction>, ModuleInstanceAddr),
     pub(crate) labels: Labels,
     pub(crate) locals: Box<[RawWasmValue]>,
 }
 
 impl CallFrame {
-    #[inline]
+    // TOOD: perf: this is called a lot, and it's a bit slow
     /// Push a new label to the label stack and ensure the stack has the correct values
     pub(crate) fn enter_label(&mut self, label_frame: LabelFrame, stack: &mut super::ValueStack) {
         if label_frame.params > 0 {
@@ -102,13 +102,15 @@ impl CallFrame {
         Some(())
     }
 
-    #[inline]
+    // TODO: perf: a lot of time is spent here
+    #[inline(always)] // about 10% faster with this
     pub(crate) fn new(
-        func_instance_ptr: Rc<FunctionInstance>,
+        wasm_func_inst: Rc<WasmFunction>,
+        owner: ModuleInstanceAddr,
         params: impl Iterator<Item = RawWasmValue> + ExactSizeIterator,
-        local_types: Vec<ValType>,
     ) -> Self {
         let locals = {
+            let local_types = &wasm_func_inst.locals;
             let total_size = local_types.len() + params.len();
             let mut locals = Vec::with_capacity(total_size);
             locals.extend(params);
@@ -116,7 +118,7 @@ impl CallFrame {
             locals.into_boxed_slice()
         };
 
-        Self { instr_ptr: 0, func_instance: func_instance_ptr, locals, labels: Labels::new() }
+        Self { instr_ptr: 0, func_instance: (wasm_func_inst, owner), locals, labels: Labels::new() }
     }
 
     #[inline]
diff --git a/crates/tinywasm/src/store.rs b/crates/tinywasm/src/store.rs
index dd8112d..a4f9b70 100644
--- a/crates/tinywasm/src/store.rs
+++ b/crates/tinywasm/src/store.rs
@@ -87,7 +87,7 @@ impl Default for Store {
 /// Data should only be addressable by the module that owns it
 /// See <https://webassembly.github.io/spec/core/exec/runtime.html#store>
 pub(crate) struct StoreData {
-    pub(crate) funcs: Vec<Rc<FunctionInstance>>,
+    pub(crate) funcs: Vec<FunctionInstance>,
     pub(crate) tables: Vec<Rc<RefCell<TableInstance>>>,
     pub(crate) memories: Vec<Rc<RefCell<MemoryInstance>>>,
     pub(crate) globals: Vec<Rc<RefCell<GlobalInstance>>>,
@@ -122,7 +122,7 @@ impl Store {
         let mut func_addrs = Vec::with_capacity(func_count);
 
         for (i, (_, func)) in funcs.into_iter().enumerate() {
-            self.data.funcs.push(Rc::new(FunctionInstance { func: Function::Wasm(func), owner: idx }));
+            self.data.funcs.push(FunctionInstance { func: Function::Wasm(Rc::new(func)), owner: idx });
             func_addrs.push((i + func_count) as FuncAddr);
         }
 
@@ -344,7 +344,7 @@ impl Store {
     }
 
     pub(crate) fn add_func(&mut self, func: Function, idx: ModuleInstanceAddr) -> Result<FuncAddr> {
-        self.data.funcs.push(Rc::new(FunctionInstance { func, owner: idx }));
+        self.data.funcs.push(FunctionInstance { func, owner: idx });
         Ok(self.data.funcs.len() as FuncAddr - 1)
     }
 
@@ -396,7 +396,7 @@ impl Store {
     }
 
     /// Get the function at the actual index in the store
-    pub(crate) fn get_func(&self, addr: usize) -> Result<&Rc<FunctionInstance>> {
+    pub(crate) fn get_func(&self, addr: usize) -> Result<&FunctionInstance> {
         self.data.funcs.get(addr).ok_or_else(|| Error::Other(format!("function {} not found", addr)))
     }
 
@@ -438,7 +438,7 @@ impl Store {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 /// A WebAssembly Function Instance
 ///
 /// See <https://webassembly.github.io/spec/core/exec/runtime.html#function-instances>
@@ -447,23 +447,6 @@ pub(crate) struct FunctionInstance {
     pub(crate) owner: ModuleInstanceAddr, // index into store.module_instances, none for host functions
 }
 
-// TODO: check if this actually helps
-#[inline(always)]
-#[cold]
-const fn cold() {}
-
-impl FunctionInstance {
-    pub(crate) fn assert_wasm(&self) -> Result<&WasmFunction> {
-        match &self.func {
-            Function::Wasm(w) => Ok(w),
-            Function::Host(_) => {
-                cold();
-                Err(Error::Other("expected wasm function".to_string()))
-            }
-        }
-    }
-}
-
 #[derive(Debug, Clone, Copy)]
 pub(crate) enum TableElement {
     Uninitialized,