chore: progress towards host functions

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

6 files changed, 92 insertions, 38 deletions

diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs
index 561d5c6..3667f1d 100644
--- a/crates/parser/src/lib.rs
+++ b/crates/parser/src/lib.rs
@@ -23,7 +23,7 @@ mod module;
 use alloc::vec::Vec;
 pub use error::*;
 use module::ModuleReader;
-use tinywasm_types::Function;
+use tinywasm_types::{Function, WasmFunction};
 use wasmparser::Validator;
 
 pub use tinywasm_types::TinyWasmModule;
@@ -108,10 +108,12 @@ impl TryFrom<ModuleReader> for TinyWasmModule {
             .code
             .into_iter()
             .zip(func_types)
-            .map(|(f, ty)| Function {
-                instructions: f.body,
-                locals: f.locals,
-                ty,
+            .map(|(f, ty)| {
+                Function::WasmFunction(WasmFunction {
+                    instructions: f.body,
+                    locals: f.locals,
+                    ty,
+                })
             })
             .collect::<Vec<_>>();
 
diff --git a/crates/tinywasm/src/func.rs b/crates/tinywasm/src/func.rs
index e422d25..ef07bf0 100644
--- a/crates/tinywasm/src/func.rs
+++ b/crates/tinywasm/src/func.rs
@@ -1,6 +1,6 @@
 use alloc::{format, string::String, string::ToString, vec, vec::Vec};
 use log::{debug, info};
-use tinywasm_types::{FuncAddr, FuncType, WasmValue};
+use tinywasm_types::{FuncAddr, FuncType, Function, WasmValue};
 
 use crate::{
     runtime::{CallFrame, Stack},
@@ -52,9 +52,11 @@ impl FuncHandle {
             }
         }
 
+        let wasm_func = &func_inst.assert_wasm()?;
+
         // 6. Let f be the dummy frame
-        debug!("locals: {:?}", func_inst.locals());
-        let call_frame = CallFrame::new(self.addr as usize, params, func_inst.locals().to_vec());
+        debug!("locals: {:?}", wasm_func.locals);
+        let call_frame = CallFrame::new(self.addr as usize, params, wasm_func.locals.to_vec());
 
         // 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)
diff --git a/crates/tinywasm/src/imports.rs b/crates/tinywasm/src/imports.rs
index 9059f29..aa9add1 100644
--- a/crates/tinywasm/src/imports.rs
+++ b/crates/tinywasm/src/imports.rs
@@ -3,7 +3,9 @@ use alloc::{
     collections::BTreeMap,
     string::{String, ToString},
 };
-use tinywasm_types::{ExternVal, ExternalKind, GlobalType, MemoryType, ModuleInstanceAddr, TableType, WasmValue};
+use tinywasm_types::{
+    ExternVal, ExternalKind, FuncAddr, GlobalType, MemoryType, ModuleInstanceAddr, TableType, WasmValue,
+};
 
 #[derive(Debug)]
 #[non_exhaustive]
@@ -19,7 +21,13 @@ pub enum Extern {
     Memory(ExternMemory),
 
     /// A function
-    Func,
+    Func(ExternFunc),
+}
+
+/// A function
+#[derive(Debug)]
+pub struct ExternFunc {
+    pub(crate) addr: FuncAddr,
 }
 
 /// A global value
@@ -69,7 +77,7 @@ impl Extern {
             Self::Global(_) => ExternalKind::Global,
             Self::Table(_) => ExternalKind::Table,
             Self::Memory(_) => ExternalKind::Memory,
-            Self::Func => ExternalKind::Func,
+            Self::Func(_) => ExternalKind::Func,
         }
     }
 }
diff --git a/crates/tinywasm/src/runtime/executor/mod.rs b/crates/tinywasm/src/runtime/executor/mod.rs
index 95a0a48..e494a27 100644
--- a/crates/tinywasm/src/runtime/executor/mod.rs
+++ b/crates/tinywasm/src/runtime/executor/mod.rs
@@ -7,7 +7,7 @@ use crate::{
     CallFrame, Error, LabelArgs, ModuleInstance, Result, Store, Trap,
 };
 use alloc::{string::ToString, vec::Vec};
-use tinywasm_types::Instruction;
+use tinywasm_types::{Function, Instruction};
 
 mod macros;
 mod traits;
@@ -25,8 +25,9 @@ impl DefaultRuntime {
         let mut cf = stack.call_stack.pop()?;
 
         // The function to execute, gets updated from ExecResult::Call
-        let mut func = store.get_func(cf.func_ptr)?.clone();
-        let mut instrs = func.instructions();
+        let mut func_inst = store.get_func(cf.func_ptr)?.clone();
+        let mut wasm_func = func_inst.assert_wasm()?;
+        let mut instrs = &wasm_func.instructions;
 
         // TODO: we might be able to index into the instructions directly
         // since the instruction pointer should always be in bounds
@@ -35,8 +36,9 @@ impl DefaultRuntime {
                 // Continue execution at the new top of the call stack
                 ExecResult::Call => {
                     cf = stack.call_stack.pop()?;
-                    func = store.get_func(cf.func_ptr)?.clone();
-                    instrs = func.instructions();
+                    func_inst = store.get_func(cf.func_ptr)?.clone();
+                    wasm_func = func_inst.assert_wasm()?;
+                    instrs = &wasm_func.instructions;
                     continue;
                 }
 
@@ -126,13 +128,19 @@ fn exec_one(
             debug!("start call");
             // prepare the call frame
             let func_idx = module.resolve_func_addr(*v);
-            let func = store.get_func(func_idx as usize)?;
-            let func_ty = module.func_ty(func.ty_addr());
+            let func_inst = store.get_func(func_idx as usize)?;
+            let func_ty = module.func_ty(func_inst.ty_addr());
 
             debug!("params: {:?}", func_ty.params);
             debug!("stack: {:?}", stack.values);
             let params = stack.values.pop_n(func_ty.params.len())?;
-            let call_frame = CallFrame::new_raw(*v as usize, &params, func.locals().to_vec());
+
+            let func = match &func_inst.func {
+                Function::WasmFunction(wasm_func) => wasm_func,
+                _ => return Err(Error::UnsupportedFeature("Host functions cannot be called".to_string())),
+            };
+
+            let call_frame = CallFrame::new_raw(*v as usize, &params, func.locals.to_vec());
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
@@ -153,8 +161,8 @@ fn exec_one(
             let func_addr = table.borrow().get(func_idx as usize)?;
 
             // prepare the call frame
-            let func = store.get_func(func_addr as usize)?;
-            let func_ty = module.func_ty(func.ty_addr());
+            let func_inst = store.get_func(func_addr as usize)?;
+            let func_ty = module.func_ty(func_inst.ty_addr());
 
             if func_ty != call_ty {
                 return Err(Trap::IndirectCallTypeMismatch {
@@ -165,7 +173,17 @@ fn exec_one(
             }
 
             let params = stack.values.pop_n(func_ty.params.len())?;
-            let call_frame = CallFrame::new_raw(func_addr as usize, &params, func.locals().to_vec());
+
+            let func = match &func_inst.func {
+                Function::WasmFunction(wasm_func) => wasm_func,
+                _ => {
+                    return Err(Error::UnsupportedFeature(
+                        "Host functions cannot be called indirectly".to_string(),
+                    ))
+                }
+            };
+
+            let call_frame = CallFrame::new_raw(func_addr as usize, &params, func.locals.to_vec());
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
diff --git a/crates/tinywasm/src/store.rs b/crates/tinywasm/src/store.rs
index 563dd8b..9c8a75e 100644
--- a/crates/tinywasm/src/store.rs
+++ b/crates/tinywasm/src/store.rs
@@ -8,7 +8,7 @@ use core::{
 use alloc::{format, rc::Rc, string::ToString, vec, vec::Vec};
 use tinywasm_types::{
     Addr, Data, Element, ElementKind, FuncAddr, Function, Global, GlobalType, Import, Instruction, MemAddr, MemoryArch,
-    MemoryType, ModuleInstanceAddr, TableAddr, TableType, TypeAddr, ValType,
+    MemoryType, ModuleInstanceAddr, TableAddr, TableType, TypeAddr, ValType, WasmFunction,
 };
 
 use crate::{
@@ -393,24 +393,27 @@ impl Store {
 /// See <https://webassembly.github.io/spec/core/exec/runtime.html#function-instances>
 pub struct FunctionInstance {
     pub(crate) func: Function,
-    pub(crate) owner: ModuleInstanceAddr, // index into store.module_instances
+    pub(crate) owner: ModuleInstanceAddr, // index into store.module_instances, none for host functions
 }
 
-impl FunctionInstance {
-    pub(crate) fn _module_instance_addr(&self) -> ModuleInstanceAddr {
-        self.owner
-    }
-
-    pub(crate) fn locals(&self) -> &[ValType] {
-        &self.func.locals
-    }
+// TODO: check if this actually helps
+#[inline(always)]
+#[cold]
+const fn cold() {}
 
-    pub(crate) fn instructions(&self) -> &[Instruction] {
-        &self.func.instructions
+impl FunctionInstance {
+    pub(crate) fn assert_wasm(&self) -> Result<&WasmFunction> {
+        match &self.func {
+            Function::WasmFunction(w) => Ok(w),
+            Function::HostFunction(_) => {
+                cold();
+                Err(Error::Other("expected wasm function".to_string()))
+            }
+        }
     }
 
     pub(crate) fn ty_addr(&self) -> TypeAddr {
-        self.func.ty
+        self.func.ty()
     }
 }
 
diff --git a/crates/types/src/lib.rs b/crates/types/src/lib.rs
index 1eef0a5..5b290ca 100644
--- a/crates/types/src/lib.rs
+++ b/crates/types/src/lib.rs
@@ -339,12 +339,33 @@ impl FuncType {
     }
 }
 
-/// A WebAssembly Function
+// A WebAssembly Function
 #[derive(Debug, Clone)]
-pub struct Function {
+pub enum Function {
+    WasmFunction(WasmFunction),
+    HostFunction(HostFunction),
+}
+
+impl Function {
+    pub fn ty(&self) -> TypeAddr {
+        match self {
+            Self::WasmFunction(f) => f.ty,
+            Self::HostFunction(f) => f.ty,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct WasmFunction {
     pub ty: TypeAddr,
-    pub locals: Box<[ValType]>,
     pub instructions: Box<[Instruction]>,
+    pub locals: Box<[ValType]>,
+}
+
+#[derive(Debug, Clone)]
+pub struct HostFunction {
+    pub ty: TypeAddr,
+    pub func: fn(&mut [WasmValue]) -> Result<(), ()>,
 }
 
 /// A WebAssembly Module Export