feat: add stateful precompile example (#8911)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>

Cargo.lock

@@ -9373,6 +9373,23 @@ version = "1.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3"
 
+[[package]]
+name = "stateful-precompile"
+version = "0.0.0"
+dependencies = [
+ "eyre",
+ "parking_lot 0.12.3",
+ "reth",
+ "reth-chainspec",
+ "reth-node-api",
+ "reth-node-core",
+ "reth-node-ethereum",
+ "reth-primitives",
+ "reth-tracing",
+ "schnellru",
+ "tokio",
+]
+
 [[package]]
 name = "static_assertions"
 version = "1.1.0"

Cargo.toml

@@ -109,6 +109,7 @@ members = [
     "examples/custom-dev-node/",
     "examples/custom-engine-types/",
     "examples/custom-evm/",
+    "examples/stateful-precompile/",
     "examples/custom-inspector/",
     "examples/custom-node-components/",
     "examples/custom-payload-builder/",

examples/stateful-precompile/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "stateful-precompile"
+version = "0.0.0"
+publish = false
+edition.workspace = true
+license.workspace = true
+
+[dependencies]
+reth.workspace = true
+reth-chainspec.workspace = true
+reth-node-api.workspace = true
+reth-node-core.workspace = true
+reth-primitives.workspace = true
+reth-node-ethereum.workspace = true
+reth-tracing.workspace = true
+
+eyre.workspace = true
+parking_lot.workspace = true
+schnellru.workspace = true
+tokio.workspace = true

examples/stateful-precompile/src/main.rs

@@ -0,0 +1,237 @@
+//! This example shows how to implement a node with a custom EVM that uses a stateful precompile
+
+#![cfg_attr(not(test), warn(unused_crate_dependencies))]
+
+use parking_lot::RwLock;
+use reth::{
+    builder::{components::ExecutorBuilder, BuilderContext, NodeBuilder},
+    primitives::{
+        revm_primitives::{CfgEnvWithHandlerCfg, Env, PrecompileResult, TxEnv},
+        Address, Bytes, U256,
+    },
+    revm::{
+        handler::register::EvmHandler,
+        inspector_handle_register,
+        precompile::{Precompile, PrecompileSpecId},
+        ContextPrecompile, ContextPrecompiles, Database, Evm, EvmBuilder, GetInspector,
+    },
+    tasks::TaskManager,
+};
+use reth_chainspec::{Chain, ChainSpec};
+use reth_node_api::{ConfigureEvm, ConfigureEvmEnv, FullNodeTypes};
+use reth_node_core::{args::RpcServerArgs, node_config::NodeConfig};
+use reth_node_ethereum::{EthEvmConfig, EthExecutorProvider, EthereumNode};
+use reth_primitives::{
+    revm_primitives::{SpecId, StatefulPrecompileMut},
+    Genesis, Header, TransactionSigned,
+};
+use reth_tracing::{RethTracer, Tracer};
+use schnellru::{ByLength, LruMap};
+use std::{collections::HashMap, sync::Arc};
+
+/// A cache for precompile inputs / outputs.
+///
+/// This assumes that the precompile is a standard precompile, as in `StandardPrecompileFn`, meaning
+/// its inputs are only `(Bytes, u64)`.
+///
+/// NOTE: This does not work with "context stateful precompiles", ie `ContextStatefulPrecompile` or
+/// `ContextStatefulPrecompileMut`. They are explicitly banned.
+#[derive(Debug, Default)]
+pub struct PrecompileCache {
+    /// Caches for each precompile input / output.
+    #[allow(clippy::type_complexity)]
+    cache: HashMap<(Address, SpecId), Arc<RwLock<LruMap<(Bytes, u64), PrecompileResult>>>>,
+}
+
+/// Custom EVM configuration
+#[derive(Debug, Clone, Default)]
+#[non_exhaustive]
+pub struct MyEvmConfig {
+    precompile_cache: Arc<RwLock<PrecompileCache>>,
+}
+
+impl MyEvmConfig {
+    /// Sets the precompiles to the EVM handler
+    ///
+    /// This will be invoked when the EVM is created via [ConfigureEvm::evm] or
+    /// [ConfigureEvm::evm_with_inspector]
+    ///
+    /// This will use the default mainnet precompiles and wrap them with a cache.
+    pub fn set_precompiles<EXT, DB>(
+        handler: &mut EvmHandler<EXT, DB>,
+        cache: Arc<RwLock<PrecompileCache>>,
+    ) where
+        DB: Database,
+    {
+        // first we need the evm spec id, which determines the precompiles
+        let spec_id = handler.cfg.spec_id;
+
+        let mut loaded_precompiles: ContextPrecompiles<DB> =
+            ContextPrecompiles::new(PrecompileSpecId::from_spec_id(spec_id));
+        for (address, precompile) in loaded_precompiles.to_mut().iter_mut() {
+            // get or insert the cache for this address / spec
+            let mut cache = cache.write();
+            let cache = cache
+                .cache
+                .entry((*address, spec_id))
+                .or_insert(Arc::new(RwLock::new(LruMap::new(ByLength::new(1024)))));
+
+            *precompile = Self::wrap_precompile(precompile.clone(), cache.clone());
+        }
+
+        // install the precompiles
+        handler.pre_execution.load_precompiles = Arc::new(move || loaded_precompiles.clone());
+    }
+
+    /// Given a [`ContextPrecompile`] and cache for a specific precompile, create a new precompile
+    /// that wraps the precompile with the cache.
+    fn wrap_precompile<DB>(
+        precompile: ContextPrecompile<DB>,
+        cache: Arc<RwLock<LruMap<(Bytes, u64), PrecompileResult>>>,
+    ) -> ContextPrecompile<DB>
+    where
+        DB: Database,
+    {
+        let ContextPrecompile::Ordinary(precompile) = precompile else {
+            // context stateful precompiles are not supported, due to lifetime issues or skill
+            // issues
+            panic!("precompile is not ordinary");
+        };
+
+        let wrapped = WrappedPrecompile { precompile, cache: cache.clone() };
+
+        ContextPrecompile::Ordinary(Precompile::StatefulMut(Box::new(wrapped)))
+    }
+}
+
+/// A custom precompile that contains the cache and precompile it wraps.
+#[derive(Clone)]
+pub struct WrappedPrecompile {
+    /// The precompile to wrap.
+    precompile: Precompile,
+    /// The cache to use.
+    cache: Arc<RwLock<LruMap<(Bytes, u64), PrecompileResult>>>,
+}
+
+impl StatefulPrecompileMut for WrappedPrecompile {
+    fn call_mut(&mut self, bytes: &Bytes, gas_price: u64, _env: &Env) -> PrecompileResult {
+        let mut cache = self.cache.write();
+        let key = (bytes.clone(), gas_price);
+
+        // get the result if it exists
+        if let Some(result) = cache.get(&key) {
+            return result.clone()
+        }
+
+        // call the precompile if cache miss
+        let output = self.precompile.call(bytes, gas_price, _env);
+        cache.insert(key, output.clone());
+
+        output
+    }
+}
+
+impl ConfigureEvmEnv for MyEvmConfig {
+    fn fill_tx_env(tx_env: &mut TxEnv, transaction: &TransactionSigned, sender: Address) {
+        EthEvmConfig::fill_tx_env(tx_env, transaction, sender)
+    }
+
+    fn fill_cfg_env(
+        cfg_env: &mut CfgEnvWithHandlerCfg,
+        chain_spec: &ChainSpec,
+        header: &Header,
+        total_difficulty: U256,
+    ) {
+        EthEvmConfig::fill_cfg_env(cfg_env, chain_spec, header, total_difficulty)
+    }
+}
+
+impl ConfigureEvm for MyEvmConfig {
+    type DefaultExternalContext<'a> = ();
+
+    fn evm<'a, DB: Database + 'a>(&self, db: DB) -> Evm<'a, Self::DefaultExternalContext<'a>, DB> {
+        let new_cache = self.precompile_cache.clone();
+        EvmBuilder::default()
+            .with_db(db)
+            // add additional precompiles
+            .append_handler_register_box(Box::new(move |handler| {
+                MyEvmConfig::set_precompiles(handler, new_cache.clone())
+            }))
+            .build()
+    }
+
+    fn evm_with_inspector<'a, DB, I>(&self, db: DB, inspector: I) -> Evm<'a, I, DB>
+    where
+        DB: Database + 'a,
+        I: GetInspector<DB>,
+    {
+        let new_cache = self.precompile_cache.clone();
+        EvmBuilder::default()
+            .with_db(db)
+            .with_external_context(inspector)
+            // add additional precompiles
+            .append_handler_register_box(Box::new(move |handler| {
+                MyEvmConfig::set_precompiles(handler, new_cache.clone())
+            }))
+            .append_handler_register(inspector_handle_register)
+            .build()
+    }
+}
+
+/// Builds a regular ethereum block executor that uses the custom EVM.
+#[derive(Debug, Default, Clone)]
+#[non_exhaustive]
+pub struct MyExecutorBuilder {
+    /// The precompile cache to use for all executors.
+    precompile_cache: Arc<RwLock<PrecompileCache>>,
+}
+
+impl<Node> ExecutorBuilder<Node> for MyExecutorBuilder
+where
+    Node: FullNodeTypes,
+{
+    type EVM = MyEvmConfig;
+    type Executor = EthExecutorProvider<Self::EVM>;
+
+    async fn build_evm(
+        self,
+        ctx: &BuilderContext<Node>,
+    ) -> eyre::Result<(Self::EVM, Self::Executor)> {
+        let evm_config = MyEvmConfig { precompile_cache: self.precompile_cache.clone() };
+        Ok((evm_config.clone(), EthExecutorProvider::new(ctx.chain_spec(), evm_config)))
+    }
+}
+
+#[tokio::main]
+async fn main() -> eyre::Result<()> {
+    let _guard = RethTracer::new().init()?;
+
+    let tasks = TaskManager::current();
+
+    // create a custom chain spec
+    let spec = ChainSpec::builder()
+        .chain(Chain::mainnet())
+        .genesis(Genesis::default())
+        .london_activated()
+        .paris_activated()
+        .shanghai_activated()
+        .cancun_activated()
+        .build();
+
+    let node_config =
+        NodeConfig::test().with_rpc(RpcServerArgs::default().with_http()).with_chain(spec);
+
+    let handle = NodeBuilder::new(node_config)
+        .testing_node(tasks.executor())
+        // configure the node with regular ethereum types
+        .with_types::<EthereumNode>()
+        // use default ethereum components but with our executor
+        .with_components(EthereumNode::components().executor(MyExecutorBuilder::default()))
+        .launch()
+        .await
+        .unwrap();
+
+    println!("Node started");
+
+    handle.node_exit_future.await
+}