tail recursive interpreter loop

Copyright (c) 2024, Arm Limited.

Signed-off-by: Karl Meakin <[email protected]>

pulley/Cargo.toml

@@ -27,7 +27,8 @@ arbitrary = ["dep:arbitrary", "arbitrary/derive", "std", "cranelift-bitset/arbit
 encode = []
 decode = []
 disas = ["decode"]
-interp = ["decode"]
+interp = ["decode", "encode"]
+tail_calls = ["interp"]
 
 [package.metadata.docs.rs]
 all-features = true

pulley/src/decode.rs

@@ -260,7 +260,8 @@ impl BytecodeStream for UnsafeBytecodeStream {
 
 /// Anything that can be decoded from a bytecode stream, e.g. opcodes,
 /// immediates, registers, etc...
-trait Decode: Sized {
+pub trait Decode: Sized {
+    /// Decode this type from the given bytecode stream.
     fn decode<T>(bytecode: &mut T) -> Result<Self, T::Error>
     where
         T: BytecodeStream;
@@ -377,6 +378,32 @@ impl Decode for PcRelOffset {
     }
 }
 
+impl Decode for Opcode {
+    fn decode<T>(bytecode: &mut T) -> Result<Self, T::Error>
+    where
+        T: BytecodeStream,
+    {
+        let byte = u8::decode(bytecode)?;
+        match Opcode::new(byte) {
+            Some(v) => Ok(v),
+            None => Err(bytecode.invalid_opcode(byte)),
+        }
+    }
+}
+
+impl Decode for ExtendedOpcode {
+    fn decode<T>(bytecode: &mut T) -> Result<Self, T::Error>
+    where
+        T: BytecodeStream,
+    {
+        let word = u16::decode(bytecode)?;
+        match ExtendedOpcode::new(word) {
+            Some(v) => Ok(v),
+            None => Err(bytecode.invalid_extended_opcode(word)),
+        }
+    }
+}
+
 impl<R: Reg> Decode for BinaryOperands<R> {
     fn decode<T>(bytecode: &mut T) -> Result<Self, T::Error>
     where
@@ -655,3 +682,43 @@ macro_rules! define_extended_decoder {
     };
 }
 for_each_extended_op!(define_extended_decoder);
+
+fn unwrap_uninhabited<T>(res: Result<T, Uninhabited>) -> T {
+    match res {
+        Ok(ok) => ok,
+    }
+}
+
+#[allow(missing_docs)]
+pub mod operands {
+    use super::*;
+
+    macro_rules! define_operands_decoder {
+        (
+            $(
+                $( #[$attr:meta] )*
+                    $snake_name:ident = $name:ident $( {
+                    $(
+                        $( #[$field_attr:meta] )*
+                        $field:ident : $field_ty:ty
+                    ),*
+                } )? ;
+            )*
+        ) => {
+            $(
+                #[allow(unused_variables)]
+                pub fn $snake_name(pc: &mut UnsafeBytecodeStream) -> ($($($field_ty,)*)?) {
+                    ($($(unwrap_uninhabited(<$field_ty>::decode(pc)),)*)?)
+                }
+            )*
+        };
+    }
+
+    for_each_op!(define_operands_decoder);
+
+    pub fn extended(pc: &mut UnsafeBytecodeStream) -> (ExtendedOpcode,) {
+        (unwrap_uninhabited(ExtendedOpcode::decode(pc)),)
+    }
+
+    for_each_extended_op!(define_operands_decoder);
+}

pulley/src/interp.rs

@@ -4,6 +4,7 @@ use crate::decode::*;
 use crate::imms::*;
 use crate::regs::*;
 use crate::ExtendedOpcode;
+use crate::Opcode;
 use alloc::string::ToString;
 use alloc::{vec, vec::Vec};
 use core::fmt;
@@ -13,6 +14,8 @@ use core::ops::{Index, IndexMut};
 use core::ptr::{self, NonNull};
 use sptr::Strict;
 
+mod interp_loop;
+
 const DEFAULT_STACK_SIZE: usize = 1 << 20; // 1 MiB
 
 /// A virtual machine for interpreting Pulley bytecode.
@@ -121,27 +124,10 @@ impl Vm {
     }
 
     unsafe fn run(&mut self, pc: NonNull<u8>) -> Result<(), NonNull<u8>> {
-        let mut visitor = InterpreterVisitor {
-            state: &mut self.state,
-            pc: UnsafeBytecodeStream::new(pc),
-        };
-
-        loop {
-            let cf = self.decoder.decode_one(&mut visitor).unwrap();
-
-            // Really wish we had `feature(explicit_tail_calls)`...
-            match cf {
-                ControlFlow::Continue(()) => continue,
-
-                // Out-of-line slow paths marked `cold` and `inline(never)` to
-                // improve codegen.
-                ControlFlow::Break(Done::Trap) => {
-                    let pc = visitor.pc.as_ptr();
-                    return self.trap(pc);
-                }
-                ControlFlow::Break(Done::ReturnToHost) => return self.return_to_host(),
-                ControlFlow::Break(Done::HostCall) => return self.host_call(),
-            }
+        match interp_loop::interpreter_loop(self, UnsafeBytecodeStream::new(pc)) {
+            Done::ReturnToHost => self.return_to_host(),
+            Done::Trap => self.trap(pc),
+            Done::HostCall => self.host_call(),
         }
     }