chore: add CallSelf/ReturnCallSelf

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

4 files changed, 65 insertions, 9 deletions

diff --git a/crates/parser/src/module.rs b/crates/parser/src/module.rs
index d498143..18bba66 100644
--- a/crates/parser/src/module.rs
+++ b/crates/parser/src/module.rs
@@ -1,11 +1,11 @@
 use crate::log::debug;
-use crate::{ParseError, Result, conversion};
+use crate::{conversion, ParseError, Result};
 use alloc::string::ToString;
 use alloc::sync::Arc;
 use alloc::{format, vec::Vec};
 use tinywasm_types::{
-    ArcSlice, Data, Element, Export, FuncType, Global, Import, Instruction, MemoryType, TableType, TinyWasmModule,
-    ValueCounts, ValueCountsSmall, WasmFunction, WasmFunctionData,
+    ArcSlice, Data, Element, Export, FuncType, Global, Import, ImportKind, Instruction, MemoryType, TableType,
+    TinyWasmModule, ValueCounts, ValueCountsSmall, WasmFunction, WasmFunctionData,
 };
 use wasmparser::{FuncValidatorAllocations, Payload, Validator};
 
@@ -31,6 +31,16 @@ pub(crate) struct ModuleReader {
 }
 
 impl ModuleReader {
+    fn apply_instruction_rewrites(instructions: &mut [Instruction], self_func_addr: u32) {
+        for instr in instructions.iter_mut() {
+            if matches!(instr, Instruction::Call(addr) if *addr == self_func_addr) {
+                *instr = Instruction::CallSelf;
+            } else if matches!(instr, Instruction::ReturnCall(addr) if *addr == self_func_addr) {
+                *instr = Instruction::ReturnCallSelf;
+            }
+        }
+    }
+
     pub(crate) fn new() -> Self {
         Self::default()
     }
@@ -189,14 +199,22 @@ impl ModuleReader {
             return Err(ParseError::Other("Code and code type address count mismatch".to_string()));
         }
 
+        let imported_func_count =
+            self.imports.iter().filter(|i| matches!(&i.kind, ImportKind::Function(_))).count() as u32;
+
         let funcs = self
             .code
             .into_iter()
             .zip(self.code_type_addrs)
-            .map(|((instructions, data, locals), ty_idx)| {
+            .enumerate()
+            .map(|(func_idx, ((instructions, data, locals), ty_idx))| {
                 let ty = self.func_types.get(ty_idx as usize).expect("No func type for func, this is a bug").clone();
                 let params = ValueCountsSmall::from(&ty.params);
-                WasmFunction { instructions: ArcSlice(instructions), data, locals, params, ty }
+                let self_func_addr = imported_func_count + func_idx as u32;
+                let mut instructions = instructions.to_vec();
+                Self::apply_instruction_rewrites(&mut instructions, self_func_addr);
+
+                WasmFunction { instructions: ArcSlice::from(instructions), data, locals, params, ty }
             })
             .collect::<Vec<_>>();
 
diff --git a/crates/tinywasm/benches/fibonacci.rs b/crates/tinywasm/benches/fibonacci.rs
index a5f1c72..eea28c4 100644
--- a/crates/tinywasm/benches/fibonacci.rs
+++ b/crates/tinywasm/benches/fibonacci.rs
@@ -36,11 +36,11 @@ fn fibonacci_run(module: TinyWasmModule, recursive: bool, n: i32) -> Result<()>
 
 fn criterion_benchmark(c: &mut Criterion) {
     let module = fibonacci_parse().expect("fibonacci_parse");
-    let twasm = fibonacci_to_twasm(&module).expect("fibonacci_to_twasm");
+    let _twasm = fibonacci_to_twasm(&module).expect("fibonacci_to_twasm");
 
-    c.bench_function("fibonacci_parse", |b| b.iter(fibonacci_parse));
-    c.bench_function("fibonacci_to_twasm", |b| b.iter(|| fibonacci_to_twasm(&module)));
-    c.bench_function("fibonacci_from_twasm", |b| b.iter(|| fibonacci_from_twasm(&twasm)));
+    // c.bench_function("fibonacci_parse", |b| b.iter(fibonacci_parse));
+    // c.bench_function("fibonacci_to_twasm", |b| b.iter(|| fibonacci_to_twasm(&module)));
+    // c.bench_function("fibonacci_from_twasm", |b| b.iter(|| fibonacci_from_twasm(&twasm)));
     c.bench_function("fibonacci_iterative_60", |b| b.iter(|| fibonacci_run(module.clone(), false, 60)));
     c.bench_function("fibonacci_recursive_26", |b| b.iter(|| fibonacci_run(module.clone(), true, 26)));
 }
diff --git a/crates/tinywasm/src/interpreter/executor.rs b/crates/tinywasm/src/interpreter/executor.rs
index d5e6fea..cb7df97 100644
--- a/crates/tinywasm/src/interpreter/executor.rs
+++ b/crates/tinywasm/src/interpreter/executor.rs
@@ -96,8 +96,10 @@ impl<'store> Executor<'store> {
             Select128 => self.store.stack.values.select::<Value128>().to_cf()?,
             SelectRef => self.store.stack.values.select::<ValueRef>().to_cf()?,
             Call(v) => return self.exec_call_direct::<false>(*v),
+            CallSelf => return self.exec_call_self::<false>(),
             CallIndirect(ty, table) => return self.exec_call_indirect::<false>(*ty, *table),
             ReturnCall(v) => return self.exec_call_direct::<true>(*v),
+            ReturnCallSelf => return self.exec_call_self::<true>(),
             ReturnCallIndirect(ty, table) => return self.exec_call_indirect::<true>(*ty, *table),
             Jump(ip) => {
                 self.cf.instr_ptr = *ip as usize;
@@ -717,6 +719,40 @@ impl<'store> Executor<'store> {
             crate::Function::Host(host_func) => self.exec_call_host(host_func.clone()),
         }
     }
+
+    fn exec_call_self<const IS_RETURN_CALL: bool>(&mut self) -> ControlFlow<Option<Error>> {
+        if !IS_RETURN_CALL && self.store.stack.call_stack.is_full() {
+            return ControlFlow::Break(Some(Trap::CallStackOverflow.into()));
+        }
+
+        let params = self.func.params;
+        let locals = self.func.locals;
+
+        if IS_RETURN_CALL {
+            self.store.stack.values.truncate_keep_counts(self.cf.locals_base, params);
+        }
+
+        let (locals_base, _stack_base, stack_offset) = match self.store.stack.values.enter_locals(params, locals) {
+            Ok(v) => v,
+            Err(Error::Trap(Trap::ValueStackOverflow)) if !IS_RETURN_CALL => {
+                return ControlFlow::Break(Some(Trap::CallStackOverflow.into()));
+            }
+            Err(err) => return ControlFlow::Break(Some(err)),
+        };
+
+        let new_call_frame = CallFrame::new(self.cf.func_addr, self.cf.module_addr, locals_base, stack_offset);
+
+        if IS_RETURN_CALL {
+            self.cf = new_call_frame;
+        } else {
+            self.cf.incr_instr_ptr();
+            self.store.stack.call_stack.push(self.cf).to_cf()?;
+            self.cf = new_call_frame;
+        }
+
+        ControlFlow::Continue(())
+    }
+
     fn exec_call_indirect<const IS_RETURN_CALL: bool>(
         &mut self,
         type_addr: u32,
diff --git a/crates/types/src/instructions.rs b/crates/types/src/instructions.rs
index b884319..2474012 100644
--- a/crates/types/src/instructions.rs
+++ b/crates/types/src/instructions.rs
@@ -67,8 +67,10 @@ pub enum Instruction {
     BranchTableTarget(u32),      // (landing_pad_ip)
     Return,
     Call(FuncAddr),
+    CallSelf,
     CallIndirect(TypeAddr, TableAddr),
     ReturnCall(FuncAddr),
+    ReturnCallSelf,
     ReturnCallIndirect(TypeAddr, TableAddr),
 
     // > Parametric Instructions