Trampolining draft

crates/interpreter/src/exec.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use alloc::boxed::Box;
 // TODO: Some toctou could be used instead of panic.
 use alloc::vec;
 use alloc::vec::Vec;
@@ -724,6 +725,7 @@ enum ThreadResult<'m> {
     Continue(Thread<'m>),
     Done(Vec<Val>),
     Host,
+    TailCall(Box<dyn FnOnce(&mut Store<'m>) -> Result<ThreadResult<'m>, Error>>),
 }
 
 impl<'m> Thread<'m> {
@@ -751,17 +753,22 @@ impl<'m> Thread<'m> {
     }
 
     fn run<'a>(mut self, store: &'a mut Store<'m>) -> Result<RunResult<'a, 'm>, Error> {
+        let mut f: Option<Box<dyn FnOnce(&mut Store<'m>) -> Result<ThreadResult<'m>, Error>>> = None;
         loop {
-            // TODO: When trapping, we could return some CoreDump<'m> that contains the Thread<'m>.
-            // This permits to dump the frames.
-            match self.step(store)? {
+            let result = if let Some(func) = f.take() {
+                func(store)?
+            } else {
+                self.step(store)?
+            };
+
+            match result {
                 ThreadResult::Continue(x) => self = x,
                 ThreadResult::Done(x) => return Ok(RunResult::Done(x)),
                 ThreadResult::Host => return Ok(RunResult::Host(Call { store })),
+                ThreadResult::TailCall(next_f) => f = Some(next_f),
             }
         }
     }
-
     fn step(mut self, store: &mut Store<'m>) -> Result<ThreadResult<'m>, Error> {
         use Instr::*;
         let saved = self.parser.save();
@@ -795,7 +802,17 @@ impl<'m> Thread<'m> {
                 return Ok(self.pop_label(inst, ls.get(i).cloned().unwrap_or(ln)));
             }
             Return => return Ok(self.exit_frame()),
-            Call(x) => return self.invoke(store, store.func_ptr(inst_id, x)),
+            Call(x) => {
+                return if self.parser.is_tail_call() {
+                    // Wrap the invoke call in a closure (trampoline)
+                    Ok(ThreadResult::TailCall(Box::new(move |store| {
+                        self.invoke(store, store.func_ptr(inst_id, x))
+                    })))
+                } else {
+                    self.invoke(store, store.func_ptr(inst_id, x))
+                }
+            }
+            // similar changes for CallIndirect
             CallIndirect(x, y) => {
                 let i = self.pop_value().unwrap_i32();
                 let x = match store.table(inst_id, x).elems.get(i as usize) {

crates/interpreter/src/parser.rs

@@ -37,6 +37,82 @@ impl<'m> Parser<'m, Use> {
     pub unsafe fn new(data: &'m [u8]) -> Self {
         Self::internal_new(data)
     }
+
+    pub fn is_tail_call(&self) -> bool {
+        let mut remaining = self.data;
+        let mut block_depth = 0;
+        let mut call_depth = 0;
+
+        while !remaining.is_empty() {
+            let mut temp_parser: Parser<'_, Check> = Parser { data: remaining, mode: PhantomData };
+
+            if remaining.len() == 0 {
+                return false;
+            } else if remaining.len() == 1 {
+                return remaining[0] == 0x0B && block_depth == 0 && call_depth == 1;
+            }
+
+            let opcode = temp_parser.parse_byte().unwrap();
+            remaining = temp_parser.data;
+
+            if opcode == 0x02 || opcode == 0x03 || opcode == 0x04 {
+                block_depth += 1;
+            } else if opcode == 0x0B {
+                block_depth -= 1;
+                if call_depth > 0 && block_depth == 0 {
+                    call_depth -= 1;
+                }
+            } else if opcode == 0x10 || opcode == 0x11 {
+                call_depth += 1;
+                if block_depth == 0 && call_depth == 1 {
+                    return !remaining.is_empty() && remaining[0] == 0x0B;
+                }
+            }
+
+            match opcode {
+                0x0E => {
+                    // br_table
+                    let _num_labels = temp_parser.parse_u32().unwrap();
+                    for _ in 0 .. _num_labels + 1 {
+                        temp_parser.parse_labelidx().unwrap();
+                    }
+                    remaining = temp_parser.data;
+                }
+                0x28 ..= 0x3E => {
+                    // Memory instructions
+                    temp_parser.parse_memarg().unwrap();
+                    remaining = temp_parser.data;
+                }
+                0xFC => {
+                    let fc_opcode = temp_parser.parse_u32().unwrap();
+                    match fc_opcode {
+                        0 ..= 3 => {
+                            // Using range pattern
+                            temp_parser.parse_leb128(true, 33).unwrap();
+                            remaining = temp_parser.data;
+                        }
+                        4 => {
+                            temp_parser.parse_dataidx().unwrap();
+                            temp_parser.parse_byte().unwrap();
+                            remaining = temp_parser.data;
+                        }
+                        5 | 6 => {
+                            temp_parser.parse_elemidx().unwrap();
+                            remaining = temp_parser.data;
+                        }
+                        7 => {
+                            temp_parser.parse_tableidx().unwrap();
+                            temp_parser.parse_elemidx().unwrap();
+                            remaining = temp_parser.data;
+                        }
+                        _ => (), // Unsupported or no arguments
+                    }
+                }
+                _ => (), // Other instructions with no immediate arguments
+            }
+        }
+        false // Not a tail call if reached end without finding one
+    }
 }
 
 impl<'m, M: Mode> Parser<'m, M> {