feat: Add boxed precompile trait (#1131)

* feat: Add boxed precompile trait

* add box

* Add mut Box and ref Arc precompiles, rename

crates/precompile/src/lib.rs

@@ -20,7 +20,7 @@ mod modexp;
 mod secp256k1;
 pub mod utilities;
 
-use core::{fmt, hash::Hash};
+use core::hash::Hash;
 use once_cell::race::OnceBox;
 #[doc(hidden)]
 pub use revm_primitives as primitives;
@@ -181,8 +181,14 @@ impl Precompiles {
 
     /// Returns the precompile for the given address.
     #[inline]
-    pub fn get(&self, address: &Address) -> Option<Precompile> {
-        self.inner.get(address).cloned()
+    pub fn get(&self, address: &Address) -> Option<&Precompile> {
+        self.inner.get(address)
+    }
+
+    /// Returns the precompile for the given address.
+    #[inline]
+    pub fn get_mut(&mut self, address: &Address) -> Option<&mut Precompile> {
+        self.inner.get_mut(address)
     }
 
     /// Is the precompiles list empty.
@@ -203,21 +209,6 @@ impl Precompiles {
     }
 }
 
-#[derive(Clone)]
-pub enum Precompile {
-    Standard(StandardPrecompileFn),
-    Env(EnvPrecompileFn),
-}
-
-impl fmt::Debug for Precompile {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match self {
-            Precompile::Standard(_) => f.write_str("Standard"),
-            Precompile::Env(_) => f.write_str("Env"),
-        }
-    }
-}
-
 #[derive(Clone, Debug)]
 pub struct PrecompileWithAddress(pub Address, pub Precompile);
 

crates/primitives/Cargo.toml

@@ -33,6 +33,7 @@ once_cell = { version = "1.19", default-features = false, optional = true }
 enumn = "0.1"
 derive_more = { version = "0.99", optional = true }
 cfg-if = "1"
+dyn-clone = "1.0"
 
 # optional
 serde = { version = "1.0", default-features = false, features = [

crates/primitives/src/precompile.rs

@@ -1,6 +1,7 @@
-use crate::Env;
-use alloy_primitives::Bytes;
+use crate::{Bytes, Env};
 use core::fmt;
+use dyn_clone::DynClone;
+use std::{boxed::Box, sync::Arc};
 
 /// A precompile operation result.
 ///
@@ -10,6 +11,98 @@ pub type PrecompileResult = Result<(u64, Bytes), PrecompileError>;
 pub type StandardPrecompileFn = fn(&Bytes, u64) -> PrecompileResult;
 pub type EnvPrecompileFn = fn(&Bytes, u64, env: &Env) -> PrecompileResult;
 
+/// Stateful precompile trait. It is used to create
+/// a arc precompile Precompile::Stateful.
+pub trait StatefulPrecompile: Sync + Send {
+    fn call(&self, bytes: &Bytes, gas_price: u64, env: &Env) -> PrecompileResult;
+}
+
+/// Mutable stateful precompile trait. It is used to create
+/// a boxed precompile in Precompile::StatefulMut.
+pub trait StatefulPrecompileMut: DynClone + Send + Sync {
+    fn call_mut(&mut self, bytes: &Bytes, gas_price: u64, env: &Env) -> PrecompileResult;
+}
+
+dyn_clone::clone_trait_object!(StatefulPrecompileMut);
+
+/// Arc over stateful precompile.
+pub type StatefulPrecompileArc = Arc<dyn StatefulPrecompile>;
+
+/// Box over mutable stateful precompile
+pub type StatefulPrecompileBox = Box<dyn StatefulPrecompileMut>;
+
+/// Precompile and its handlers.
+#[derive(Clone)]
+pub enum Precompile {
+    /// Standard simple precompile that takes input and gas limit.
+    Standard(StandardPrecompileFn),
+    /// Similar to Standard but takes reference to environment.
+    Env(EnvPrecompileFn),
+    /// Stateful precompile that is Arc over [`StatefulPrecompile`] trait.
+    /// It takes a reference to input, gas limit and environment.
+    Stateful(StatefulPrecompileArc),
+    /// Mutable stateful precompile that is Box over [`StatefulPrecompileMut`] trait.
+    /// It takes a reference to input, gas limit and environment.
+    StatefulMut(StatefulPrecompileBox),
+}
+
+impl From<StandardPrecompileFn> for Precompile {
+    fn from(p: StandardPrecompileFn) -> Self {
+        Precompile::Standard(p)
+    }
+}
+
+impl From<EnvPrecompileFn> for Precompile {
+    fn from(p: EnvPrecompileFn) -> Self {
+        Precompile::Env(p)
+    }
+}
+
+impl From<StatefulPrecompileArc> for Precompile {
+    fn from(p: StatefulPrecompileArc) -> Self {
+        Precompile::Stateful(p)
+    }
+}
+
+impl From<StatefulPrecompileBox> for Precompile {
+    fn from(p: StatefulPrecompileBox) -> Self {
+        Precompile::StatefulMut(p)
+    }
+}
+
+impl fmt::Debug for Precompile {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Precompile::Standard(_) => f.write_str("Standard"),
+            Precompile::Env(_) => f.write_str("Env"),
+            Precompile::Stateful(_) => f.write_str("Stateful"),
+            Precompile::StatefulMut(_) => f.write_str("StatefulMut"),
+        }
+    }
+}
+
+impl Precompile {
+    /// Create a new stateful precompile.
+    pub fn new_stateful<P: StatefulPrecompile + 'static>(p: P) -> Self {
+        Self::Stateful(Arc::new(p))
+    }
+
+    /// Create a new mutable stateful precompile.
+    pub fn new_stateful_mut<P: StatefulPrecompileMut + 'static>(p: P) -> Self {
+        Self::StatefulMut(Box::new(p))
+    }
+
+    /// Call the precompile with the given input and gas limit and return the result.
+    pub fn call(&mut self, bytes: &Bytes, gas_price: u64, env: &Env) -> PrecompileResult {
+        match self {
+            Precompile::Standard(p) => p(bytes, gas_price),
+            Precompile::Env(p) => p(bytes, gas_price, env),
+            Precompile::Stateful(p) => p.call(bytes, gas_price, env),
+            Precompile::StatefulMut(p) => p.call_mut(bytes, gas_price, env),
+        }
+    }
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub enum PrecompileError {
     /// out of gas is the main error. Others are here just for completeness
@@ -63,3 +156,33 @@ impl fmt::Display for PrecompileError {
         }
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn stateful_precompile_mut() {
+        #[derive(Default, Clone)]
+        struct MyPrecompile {}
+
+        impl StatefulPrecompileMut for MyPrecompile {
+            fn call_mut(
+                &mut self,
+                _bytes: &Bytes,
+                _gas_price: u64,
+                _env: &Env,
+            ) -> PrecompileResult {
+                PrecompileResult::Err(PrecompileError::OutOfGas)
+            }
+        }
+
+        let mut p = Precompile::new_stateful_mut(MyPrecompile::default());
+        match &mut p {
+            Precompile::StatefulMut(p) => {
+                let _ = p.call_mut(&Bytes::new(), 0, &Env::default());
+            }
+            _ => panic!("not a state"),
+        }
+    }
+}

crates/revm/src/context.rs

@@ -411,8 +411,8 @@ impl<DB: Database> EvmContext<DB> {
             return return_result(result);
         }
 
-        if let Some(precompile) = self.precompiles.get(&inputs.contract) {
-            let result = self.call_precompile(precompile, inputs, gas);
+        if let Some(precompile) = self.precompiles.get_mut(&inputs.contract) {
+            let result = Self::call_precompile(precompile, &inputs.input, gas, &self.env);
             if matches!(result.result, return_ok!()) {
                 self.journaled_state.checkpoint_commit();
             } else {
@@ -444,16 +444,12 @@ impl<DB: Database> EvmContext<DB> {
     /// Call precompile contract
     #[inline]
     fn call_precompile(
-        &mut self,
-        precompile: Precompile,
-        inputs: &CallInputs,
+        precompile: &mut Precompile,
+        input_data: &Bytes,
         gas: Gas,
+        env: &Env,
     ) -> InterpreterResult {
-        let input_data = &inputs.input;
-        let out = match precompile {
-            Precompile::Standard(fun) => fun(input_data, gas.limit()),
-            Precompile::Env(fun) => fun(input_data, gas.limit(), self.env()),
-        };
+        let out = precompile.call(input_data, gas.limit(), env);
 
         let mut result = InterpreterResult {
             result: InstructionResult::Return,