feat: debug_assertions/panics, function call hooks, gas in tests, ben…

…chmarks

crates/revm-jit-core/src/context.rs

@@ -1,47 +1,250 @@
-use revm_interpreter::InstructionResult;
+use revm_interpreter::{Gas, InstructionResult};
 use revm_primitives::U256;
-use std::mem::MaybeUninit;
+use std::{fmt, mem::MaybeUninit, ptr};
 
-/// The signature of a JIT'd EVM bytecode.
-pub type JitEvmFn =
-    unsafe extern "C" fn(stack: *mut ContextStack, gas_limit: usize) -> InstructionResult;
+/// The raw function signature of a JIT'd EVM bytecode.
+///
+/// Prefer using [`JitEvmFn`] instead of this type. See [`JitEvmFn::call`] for more information.
+pub type RawJitEvmFn = unsafe extern "C" fn(
+    gas: *mut Gas,
+    stack: *mut EvmStack,
+    stack_len: *mut usize,
+) -> InstructionResult;
+
+/// A JIT'd EVM bytecode function.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct JitEvmFn(RawJitEvmFn);
+
+impl JitEvmFn {
+    /// Wraps the function.
+    #[inline]
+    pub const fn new(f: RawJitEvmFn) -> Self {
+        Self(f)
+    }
+
+    /// Unwraps the function.
+    #[inline]
+    pub const fn into_inner(self) -> RawJitEvmFn {
+        self.0
+    }
+
+    /// Calls the function.
+    ///
+    /// Arguments:
+    /// - `stack`: Pointer to the stack. Must be `Some` if `pass_stack_through_args` is set to
+    ///   `true`.
+    /// - `stack_len`: Pointer to the stack length. Must be `Some` if `pass_stack_len_through_args`
+    ///   is set to `true`.
+    /// - `gas`: Pointer to the gas object. Must be `Some` if `disable_gas` is set to `false` (the
+    ///   default).
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that the arguments are valid.
+    #[inline(always)]
+    pub unsafe fn call(
+        self,
+        gas: Option<&mut Gas>,
+        stack: Option<&mut EvmStack>,
+        stack_len: Option<&mut usize>,
+    ) -> InstructionResult {
+        (self.0)(option_as_mut_ptr(gas), option_as_mut_ptr(stack), option_as_mut_ptr(stack_len))
+    }
+}
 
 /// JIT EVM context stack.
-#[repr(C, align(32))]
+#[repr(C)]
 #[allow(missing_debug_implementations)]
-pub struct ContextStack([MaybeUninit<u8>; 32 * 1024]);
+pub struct EvmStack([MaybeUninit<EvmWord>; 1024]);
 
 #[allow(clippy::new_without_default)]
-impl ContextStack {
-    /// The size of the stack.
-    pub const SIZE: usize = 32 * 1024;
+impl EvmStack {
+    /// The size of the stack in bytes.
+    pub const SIZE: usize = 32 * Self::CAPACITY;
+
+    /// The size of the stack in U256 elements.
+    pub const CAPACITY: usize = 1024;
 
-    /// Creates a new stack.
+    /// Creates a new EVM stack, allocated on the stack.
+    ///
+    /// Use [`EvmStack::create_vec`] to create a stack on the heap.
     #[inline]
     pub fn new() -> Self {
         Self(unsafe { MaybeUninit::uninit().assume_init() })
     }
 
-    /// Returns the stack as a slice.
+    /// Creates a vector that can be used as a stack.
     #[inline]
-    pub fn as_slice(&self) -> &[u8; ContextStack::SIZE] {
-        unsafe { std::mem::transmute(&self.0) }
+    pub fn new_heap() -> Vec<EvmWord> {
+        Vec::with_capacity(1024)
+    }
+
+    /// Creates a stack from a mutable vector.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the vector's capacity is less than the stack size.
+    #[inline]
+    pub fn from_vec(vec: &Vec<EvmWord>) -> &Self {
+        assert!(vec.capacity() >= Self::CAPACITY);
+        unsafe { &*vec.as_ptr().cast() }
+    }
+
+    /// Creates a stack from a mutable vector.
+    ///
+    /// The JIT'd function will overwrite the internal contents of the vector, and will not
+    /// set the length. This is simply to have the stack allocated on the heap.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the vector's capacity is less than the stack size.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use revm_jit_core::EvmStack;
+    /// let mut stack_buf = EvmStack::new_heap();
+    /// let stack = EvmStack::from_mut_vec(&mut stack_buf);
+    /// assert_eq!(stack.as_slice().len(), EvmStack::CAPACITY);
+    /// ```
+    #[inline]
+    pub fn from_mut_vec(vec: &mut Vec<EvmWord>) -> &mut Self {
+        assert!(vec.capacity() >= Self::CAPACITY);
+        unsafe { &mut *vec.as_mut_ptr().cast() }
+    }
+
+    /// Returns the stack as a byte array.
+    #[inline]
+    pub const fn as_bytes(&self) -> &[u8; EvmStack::SIZE] {
+        unsafe { &*self.0.as_ptr().cast() }
+    }
+
+    /// Returns the stack as a byte array.
+    #[inline]
+    pub fn as_bytes_mut(&mut self) -> &mut [u8; EvmStack::SIZE] {
+        unsafe { &mut *self.0.as_mut_ptr().cast() }
     }
 
     /// Returns the stack as a slice.
     #[inline]
-    pub fn as_mut_slice(&mut self) -> &mut [u8; ContextStack::SIZE] {
-        unsafe { std::mem::transmute(&mut self.0) }
+    pub const fn as_slice(&self) -> &[EvmWord; EvmStack::CAPACITY] {
+        unsafe { &*self.0.as_ptr().cast() }
     }
 
-    /// Returns the word at the given index.
-    pub fn word(&self, index: usize) -> U256 {
-        let offset = index * 32;
-        let bytes = &self.as_slice()[offset..offset + 32];
-        if cfg!(target_endian = "big") {
-            U256::from_be_slice(bytes)
-        } else {
-            U256::from_le_slice(bytes)
+    /// Returns the stack as a mutable slice.
+    #[inline]
+    pub fn as_mut_slice(&mut self) -> &mut [EvmWord; EvmStack::CAPACITY] {
+        unsafe { &mut *self.0.as_mut_ptr().cast() }
+    }
+}
+
+/// A native 256-bit unsigned integer. This is aligned to 32 bytes.
+#[repr(C, align(32))]
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub struct EvmWord([u8; 32]);
+
+macro_rules! fmt_impl {
+    ($trait:ident) => {
+        impl fmt::$trait for EvmWord {
+            #[inline]
+            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+                self.to_u256().fmt(f)
+            }
         }
+    };
+}
+
+fmt_impl!(Debug);
+fmt_impl!(Display);
+fmt_impl!(LowerHex);
+fmt_impl!(UpperHex);
+fmt_impl!(Binary);
+
+impl From<U256> for EvmWord {
+    #[inline]
+    fn from(value: U256) -> Self {
+        #[cfg(target_endian = "little")]
+        return unsafe { std::mem::transmute(value) };
+        #[cfg(target_endian = "big")]
+        return Self(value.to_be_bytes());
+    }
+}
+
+impl From<&U256> for EvmWord {
+    #[inline]
+    fn from(value: &U256) -> Self {
+        Self::from(*value)
+    }
+}
+
+impl From<&mut U256> for EvmWord {
+    #[inline]
+    fn from(value: &mut U256) -> Self {
+        Self::from(*value)
+    }
+}
+
+impl EvmWord {
+    /// The zero value.
+    pub const ZERO: Self = Self([0; 32]);
+
+    /// Creates a new value.
+    #[inline]
+    pub const fn new(x: [u8; 32]) -> Self {
+        Self(x)
+    }
+
+    /// Converts a [`U256`].
+    #[inline]
+    pub const fn from_u256(value: U256) -> Self {
+        #[cfg(target_endian = "little")]
+        return unsafe { std::mem::transmute(value) };
+        #[cfg(target_endian = "big")]
+        return Self(value.to_be_bytes());
+    }
+
+    /// Converts a [`U256`] reference to a [`U256`].
+    #[inline]
+    #[cfg(target_endian = "little")]
+    pub const fn from_u256_ref(value: &U256) -> &Self {
+        unsafe { &*(value as *const U256 as *const Self) }
+    }
+
+    /// Converts a [`U256`] mutable reference to a [`U256`].
+    #[inline]
+    #[cfg(target_endian = "little")]
+    pub fn from_u256_mut(value: &mut U256) -> &mut Self {
+        unsafe { &mut *(value as *mut U256 as *mut Self) }
+    }
+
+    /// Casts this value to a [`U256`]. This is a no-op on little-endian systems.
+    #[cfg(target_endian = "little")]
+    #[inline]
+    pub const fn as_u256(&self) -> &U256 {
+        unsafe { &*(self as *const Self as *const U256) }
+    }
+
+    /// Casts this value to a [`U256`]. This is a no-op on little-endian systems.
+    #[cfg(target_endian = "little")]
+    #[inline]
+    pub fn as_u256_mut(&mut self) -> &mut U256 {
+        unsafe { &mut *(self as *mut Self as *mut U256) }
+    }
+
+    /// Converts this value to a [`U256`]. This is a simple copy on little-endian systems.
+    #[inline]
+    pub const fn to_u256(&self) -> U256 {
+        #[cfg(target_endian = "little")]
+        return *self.as_u256();
+        #[cfg(target_endian = "big")]
+        return U256::from_be_bytes(self.0);
+    }
+}
+
+#[inline(always)]
+fn option_as_mut_ptr<T>(opt: Option<&mut T>) -> *mut T {
+    match opt {
+        Some(ref_) => ref_,
+        None => ptr::null_mut(),
     }
 }

crates/revm-jit-core/src/traits.rs

@@ -1,4 +1,4 @@
-use crate::{JitEvmFn, Result};
+use crate::{RawJitEvmFn, Result};
 use revm_primitives::U256;
 use std::{fmt, path::Path};
 
@@ -40,17 +40,22 @@ pub enum OptimizationLevel {
     Aggressive,
 }
 
-pub trait Builder {
+pub trait BuilderTypes {
     type Type: Copy + Eq + fmt::Debug;
     type Value: Copy + Eq + fmt::Debug;
     type StackSlot: Copy + Eq + fmt::Debug;
     type BasicBlock: Copy + Eq + fmt::Debug;
+    type Function: Copy + Eq + fmt::Debug;
+}
 
+pub trait TypeMethods: BuilderTypes {
     fn type_ptr(&self) -> Self::Type;
     fn type_ptr_sized_int(&self) -> Self::Type;
     fn type_int(&self, bits: u32) -> Self::Type;
     fn type_array(&self, ty: Self::Type, size: u32) -> Self::Type;
+}
 
+pub trait Builder: BuilderTypes + TypeMethods {
     fn create_block(&mut self, name: &str) -> Self::BasicBlock;
     fn create_block_after(&mut self, after: Self::BasicBlock, name: &str) -> Self::BasicBlock;
     fn switch_to_block(&mut self, block: Self::BasicBlock);
@@ -69,6 +74,7 @@ pub trait Builder {
     fn new_stack_slot(&mut self, ty: Self::Type, name: &str) -> Self::StackSlot;
     fn stack_load(&mut self, ty: Self::Type, slot: Self::StackSlot, name: &str) -> Self::Value;
     fn stack_store(&mut self, value: Self::Value, slot: Self::StackSlot);
+    fn stack_addr(&mut self, stack_slot: Self::StackSlot) -> Self::Value;
 
     fn load(&mut self, ty: Self::Type, ptr: Self::Value, name: &str) -> Self::Value;
     fn store(&mut self, value: Self::Value, ptr: Self::Value);
@@ -78,6 +84,8 @@ pub trait Builder {
 
     fn icmp(&mut self, cond: IntCC, lhs: Self::Value, rhs: Self::Value) -> Self::Value;
     fn icmp_imm(&mut self, cond: IntCC, lhs: Self::Value, rhs: i64) -> Self::Value;
+    fn is_null(&mut self, ptr: Self::Value) -> Self::Value;
+    fn is_not_null(&mut self, ptr: Self::Value) -> Self::Value;
     fn br(&mut self, dest: Self::BasicBlock);
     fn brif(
         &mut self,
@@ -125,22 +133,39 @@ pub trait Builder {
     fn sext(&mut self, ty: Self::Type, value: Self::Value) -> Self::Value;
 
     fn gep(&mut self, ty: Self::Type, ptr: Self::Value, offset: Self::Value) -> Self::Value;
+
+    fn panic(&mut self, msg: &str);
 }
 
-pub trait Backend {
-    type Builder<'a>: Builder
+pub trait Backend: BuilderTypes + TypeMethods {
+    type Builder<'a>: Builder<
+        Type = Self::Type,
+        Value = Self::Value,
+        StackSlot = Self::StackSlot,
+        BasicBlock = Self::BasicBlock,
+        Function = Self::Function,
+    >
     where
         Self: 'a;
 
     fn ir_extension(&self) -> &'static str;
 
     fn set_is_dumping(&mut self, yes: bool);
+    fn set_debug_assertions(&mut self, yes: bool);
     fn set_opt_level(&mut self, level: OptimizationLevel);
     fn dump_ir(&mut self, path: &Path) -> Result<()>;
     fn dump_disasm(&mut self, path: &Path) -> Result<()>;
 
     fn build_function(&mut self, name: &str) -> Result<Self::Builder<'_>>;
     fn verify_function(&mut self, name: &str) -> Result<()>;
     fn optimize_function(&mut self, name: &str) -> Result<()>;
-    fn get_function(&mut self, name: &str) -> Result<JitEvmFn>;
+    fn get_function(&mut self, name: &str) -> Result<RawJitEvmFn>;
+
+    fn add_callback_function(
+        &mut self,
+        name: &str,
+        ret: Self::Type,
+        params: &[Self::Type],
+        address: usize,
+    ) -> Self::Function;
 }