feat(evm): MultiFork generic over BlockEnv (foundry-rs#14033)

mablr · web-flow · commit b9210c196bb8 · 2026-03-31T20:14:02.000+02:00
diff --git a/crates/evm/core/src/backend/mod.rs b/crates/evm/core/src/backend/mod.rs
@@ -441,7 +441,7 @@ struct _ObjectSafe(dyn DatabaseExt);
 #[must_use]
 pub struct Backend<N: Network = AnyNetwork> {
     /// The access point for managing forks
-    forks: MultiFork<N, SpecId>,
+    forks: MultiFork<N, SpecId, BlockEnv>,
     // The default in memory db
     mem_db: FoundryEvmInMemoryDB,
     /// The journaled_state to use to initialize new forks with
@@ -478,7 +478,7 @@ where
     /// If `fork` is `Some` this will use a `fork` database, otherwise with an in-memory
     /// database.
     pub fn spawn(fork: Option<CreateFork>) -> eyre::Result<Self> {
-        Self::new(MultiFork::<N, SpecId>::spawn(), fork)
+        Self::new(MultiFork::<N, SpecId, BlockEnv>::spawn(), fork)
     }
 
     /// Creates a new instance of `Backend`
@@ -487,7 +487,10 @@ where
     /// database.
     ///
     /// Prefer using [`spawn`](Self::spawn) instead.
-    pub fn new(forks: MultiFork<N, SpecId>, fork: Option<CreateFork>) -> eyre::Result<Self> {
+    pub fn new(
+        forks: MultiFork<N, SpecId, BlockEnv>,
+        fork: Option<CreateFork>,
+    ) -> eyre::Result<Self> {
         trace!(target: "backend", forking_mode=?fork.is_some(), "creating executor backend");
         // Note: this will take of registering the `fork`
         let inner = BackendInner {
diff --git a/crates/evm/core/src/fork/multi.rs b/crates/evm/core/src/fork/multi.rs
@@ -4,18 +4,21 @@
 //! concurrently active pairs at once.
 
 use super::CreateFork;
+use crate::FoundryBlock;
 use alloy_evm::EvmEnv;
 use alloy_network::Network;
 use alloy_primitives::{U256, map::HashMap};
 use foundry_config::Config;
-use foundry_fork_db::{BackendHandler, BlockchainDb, SharedBackend, cache::BlockchainDbMeta};
+use foundry_fork_db::{
+    BackendHandler, BlockchainDb, ForkBlockEnv, SharedBackend, cache::BlockchainDbMeta,
+};
 use futures::{
     FutureExt, StreamExt,
     channel::mpsc::{Receiver, Sender, channel},
     stream::Fuse,
     task::{Context, Poll},
 };
-use revm::{context::BlockEnv, primitives::hardfork::SpecId};
+use revm::primitives::hardfork::SpecId;
 use std::{
     fmt::{self, Write},
     pin::Pin,
@@ -66,15 +69,18 @@ impl<T: Into<String>> From<T> for ForkId {
 /// Can send requests to the `MultiForkHandler` to create forks.
 #[derive(Clone, Debug)]
 #[must_use]
-pub struct MultiFork<N: Network, SPEC, BLOCK = BlockEnv> {
+pub struct MultiFork<N: Network, SPEC, BLOCK: ForkBlockEnv> {
     /// Channel to send `Request`s to the handler.
     handler: Sender<Request<N, SPEC, BLOCK>>,
     /// Ensures that all rpc resources get flushed properly.
     _shutdown: Arc<ShutDownMultiFork<N, SPEC, BLOCK>>,
 }
 
-impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + Unpin + Send + 'static>
-    MultiFork<N, SPEC>
+impl<
+    N: Network,
+    SPEC: Into<SpecId> + Default + Copy + Clone + Unpin + Send + 'static,
+    BLOCK: FoundryBlock + ForkBlockEnv + Default + Unpin,
+> MultiFork<N, SPEC, BLOCK>
 {
     /// Creates a new pair and spawns the `MultiForkHandler` on a background thread.
     pub fn spawn() -> Self {
@@ -117,7 +123,7 @@ impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + Unpin + Send + 's
     ///
     /// Use [`spawn`](Self::spawn) instead.
     #[doc(hidden)]
-    pub fn new() -> (Self, MultiForkHandler<N, SPEC>) {
+    pub fn new() -> (Self, MultiForkHandler<N, SPEC, BLOCK>) {
         let (handler, handler_rx) = channel(1);
         let _shutdown = Arc::new(ShutDownMultiFork { handler: Some(handler.clone()) });
         (Self { handler, _shutdown }, MultiForkHandler::new(handler_rx))
@@ -126,10 +132,11 @@ impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + Unpin + Send + 's
     /// Returns a fork backend.
     ///
     /// If no matching fork backend exists it will be created.
+    #[allow(clippy::type_complexity)]
     pub fn create_fork(
         &self,
         fork: CreateFork,
-    ) -> eyre::Result<(ForkId, SharedBackend<N>, EvmEnv<SPEC>)> {
+    ) -> eyre::Result<(ForkId, SharedBackend<N, BLOCK>, EvmEnv<SPEC, BLOCK>)> {
         trace!("Creating new fork, url={}, block={:?}", fork.url, fork.evm_opts.fork_block_number);
         let (sender, rx) = oneshot_channel();
         let req = Request::CreateFork(Box::new(fork), sender);
@@ -140,11 +147,12 @@ impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + Unpin + Send + 's
     /// Rolls the block of the fork.
     ///
     /// If no matching fork backend exists it will be created.
+    #[allow(clippy::type_complexity)]
     pub fn roll_fork(
         &self,
         fork: ForkId,
         block: u64,
-    ) -> eyre::Result<(ForkId, SharedBackend<N>, EvmEnv<SPEC>)> {
+    ) -> eyre::Result<(ForkId, SharedBackend<N, BLOCK>, EvmEnv<SPEC, BLOCK>)> {
         trace!(?fork, ?block, "rolling fork");
         let (sender, rx) = oneshot_channel();
         let req = Request::RollFork(fork, block, sender);
@@ -153,7 +161,7 @@ impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + Unpin + Send + 's
     }
 
     /// Returns the `EvmEnv` of the given fork, if any.
-    pub fn get_evm_env(&self, fork: ForkId) -> eyre::Result<Option<EvmEnv<SPEC>>> {
+    pub fn get_evm_env(&self, fork: ForkId) -> eyre::Result<Option<EvmEnv<SPEC, BLOCK>>> {
         trace!(?fork, "getting env config");
         let (sender, rx) = oneshot_channel();
         let req = Request::GetEvmEnv(fork, sender);
@@ -174,7 +182,10 @@ impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + Unpin + Send + 's
     ///
     /// This is required for tx level forking where we need to fork off the `block - 1` state but
     /// still need use env settings for `env`.
-    pub fn update_block_env(&self, fork: ForkId, env: BlockEnv) -> eyre::Result<()> {
+    pub fn update_block_env(&self, fork: ForkId, env: BLOCK) -> eyre::Result<()>
+    where
+        BLOCK: fmt::Debug,
+    {
         trace!(?fork, ?env, "update fork block");
         self.handler
             .clone()
@@ -185,7 +196,7 @@ impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + Unpin + Send + 's
     /// Returns the corresponding fork if it exists.
     ///
     /// Returns `None` if no matching fork backend is available.
-    pub fn get_fork(&self, id: impl Into<ForkId>) -> eyre::Result<Option<SharedBackend<N>>> {
+    pub fn get_fork(&self, id: impl Into<ForkId>) -> eyre::Result<Option<SharedBackend<N, BLOCK>>> {
         let id = id.into();
         trace!(?id, "get fork backend");
         let (sender, rx) = oneshot_channel();
@@ -207,21 +218,26 @@ impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + Unpin + Send + 's
 
 type CreateFuture<N, SPEC, BLOCK> = Pin<
     Box<
-        dyn Future<Output = eyre::Result<(ForkId, CreatedFork<N, SPEC, BLOCK>, BackendHandler<N>)>>
-            + Send,
+        dyn Future<
+                Output = eyre::Result<(
+                    ForkId,
+                    CreatedFork<N, SPEC, BLOCK>,
+                    BackendHandler<N, BLOCK>,
+                )>,
+            > + Send,
     >,
 >;
 type CreateSender<N, SPEC, BLOCK> =
-    OneshotSender<eyre::Result<(ForkId, SharedBackend<N>, EvmEnv<SPEC, BLOCK>)>>;
+    OneshotSender<eyre::Result<(ForkId, SharedBackend<N, BLOCK>, EvmEnv<SPEC, BLOCK>)>>;
 type GetEvmEnvSender<SPEC, BLOCK> = OneshotSender<Option<EvmEnv<SPEC, BLOCK>>>;
 
 /// Request that's send to the handler.
 #[derive(Debug)]
-enum Request<N: Network, SPEC, BLOCK = BlockEnv> {
+enum Request<N: Network, SPEC, BLOCK: ForkBlockEnv> {
     /// Creates a new ForkBackend.
     CreateFork(Box<CreateFork>, CreateSender<N, SPEC, BLOCK>),
     /// Returns the Fork backend for the `ForkId` if it exists.
-    GetFork(ForkId, OneshotSender<Option<SharedBackend<N>>>),
+    GetFork(ForkId, OneshotSender<Option<SharedBackend<N, BLOCK>>>),
     /// Adjusts the block that's being forked, by creating a new fork at the new block.
     RollFork(ForkId, u64, CreateSender<N, SPEC, BLOCK>),
     /// Returns the environment of the fork.
@@ -236,7 +252,7 @@ enum Request<N: Network, SPEC, BLOCK = BlockEnv> {
     GetForkUrl(ForkId, OneshotSender<Option<String>>),
 }
 
-enum ForkTask<N: Network, SPEC, BLOCK = BlockEnv> {
+enum ForkTask<N: Network, SPEC, BLOCK: ForkBlockEnv> {
     /// Contains the future that will establish a new fork.
     Create(
         CreateFuture<N, SPEC, BLOCK>,
@@ -248,14 +264,14 @@ enum ForkTask<N: Network, SPEC, BLOCK = BlockEnv> {
 
 /// The type that manages connections in the background.
 #[must_use = "futures do nothing unless polled"]
-pub struct MultiForkHandler<N: Network, SPEC, BLOCK = BlockEnv> {
+pub struct MultiForkHandler<N: Network, SPEC, BLOCK: ForkBlockEnv> {
     /// Incoming requests from the `MultiFork`.
     incoming: Fuse<Receiver<Request<N, SPEC, BLOCK>>>,
 
     /// All active handlers.
     ///
     /// It's expected that this list will be rather small (<10).
-    handlers: Vec<(ForkId, BackendHandler<N>)>,
+    handlers: Vec<(ForkId, BackendHandler<N, BLOCK>)>,
 
     // tasks currently in progress
     pending_tasks: Vec<ForkTask<N, SPEC, BLOCK>>,
@@ -270,8 +286,13 @@ pub struct MultiForkHandler<N: Network, SPEC, BLOCK = BlockEnv> {
     flush_cache_interval: Option<tokio::time::Interval>,
 }
 
-impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + 'static> MultiForkHandler<N, SPEC> {
-    fn new(incoming: Receiver<Request<N, SPEC>>) -> Self {
+impl<
+    N: Network,
+    SPEC: Into<SpecId> + Default + Copy + Clone + 'static,
+    BLOCK: FoundryBlock + ForkBlockEnv + Default,
+> MultiForkHandler<N, SPEC, BLOCK>
+{
+    fn new(incoming: Receiver<Request<N, SPEC, BLOCK>>) -> Self {
         Self {
             incoming: incoming.fuse(),
             handlers: Default::default(),
@@ -292,7 +313,7 @@ impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + 'static> MultiFor
     fn find_in_progress_task(
         &mut self,
         id: &ForkId,
-    ) -> Option<&mut Vec<CreateSender<N, SPEC, BlockEnv>>> {
+    ) -> Option<&mut Vec<CreateSender<N, SPEC, BLOCK>>> {
         for ForkTask::Create(_, in_progress, _, additional) in &mut self.pending_tasks {
             if in_progress == id {
                 return Some(additional);
@@ -301,7 +322,7 @@ impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + 'static> MultiFor
         None
     }
 
-    fn create_fork(&mut self, fork: CreateFork, sender: CreateSender<N, SPEC, BlockEnv>) {
+    fn create_fork(&mut self, fork: CreateFork, sender: CreateSender<N, SPEC, BLOCK>) {
         let fork_id = ForkId::new(&fork.url, fork.evm_opts.fork_block_number);
         trace!(?fork_id, "created new forkId");
 
@@ -319,9 +340,9 @@ impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + 'static> MultiFor
     fn insert_new_fork(
         &mut self,
         fork_id: ForkId,
-        fork: CreatedFork<N, SPEC>,
-        sender: CreateSender<N, SPEC, BlockEnv>,
-        additional_senders: Vec<CreateSender<N, SPEC, BlockEnv>>,
+        fork: CreatedFork<N, SPEC, BLOCK>,
+        sender: CreateSender<N, SPEC, BLOCK>,
+        additional_senders: Vec<CreateSender<N, SPEC, BLOCK>>,
     ) {
         self.forks.insert(fork_id.clone(), fork.clone());
         let _ = sender.send(Ok((fork_id.clone(), fork.backend.clone(), fork.evm_env.clone())));
@@ -335,7 +356,7 @@ impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + 'static> MultiFor
     }
 
     /// Update the fork's block entire env
-    fn update_env(&mut self, fork_id: ForkId, env: BlockEnv) {
+    fn update_env(&mut self, fork_id: ForkId, env: BLOCK) {
         if let Some(fork) = self.forks.get_mut(&fork_id) {
             fork.evm_env.block_env = env;
         }
@@ -344,12 +365,12 @@ impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + 'static> MultiFor
     /// cheatcodes when new fork selected.
     fn update_block(&mut self, fork_id: ForkId, block_number: U256, block_timestamp: U256) {
         if let Some(fork) = self.forks.get_mut(&fork_id) {
-            fork.evm_env.block_env.number = block_number;
-            fork.evm_env.block_env.timestamp = block_timestamp;
+            fork.evm_env.block_env.set_number(block_number);
+            fork.evm_env.block_env.set_timestamp(block_timestamp);
         }
     }
 
-    fn on_request(&mut self, req: Request<N, SPEC>) {
+    fn on_request(&mut self, req: Request<N, SPEC, BLOCK>) {
         match req {
             Request::CreateFork(fork, sender) => self.create_fork(*fork, sender),
             Request::GetFork(fork_id, sender) => {
@@ -393,8 +414,11 @@ impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + 'static> MultiFor
 
 // Drives all handler to completion.
 // This future will finish once all underlying BackendHandler are completed.
-impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + Unpin + 'static> Future
-    for MultiForkHandler<N, SPEC>
+impl<
+    N: Network,
+    SPEC: Into<SpecId> + Default + Copy + Clone + Unpin + 'static,
+    BLOCK: FoundryBlock + ForkBlockEnv + Default + Unpin,
+> Future for MultiForkHandler<N, SPEC, BLOCK>
 {
     type Output = ();
 
@@ -496,20 +520,24 @@ impl<N: Network, SPEC: Into<SpecId> + Default + Copy + Clone + Unpin + 'static>
 
 /// Tracks the created Fork
 #[derive(Debug, Clone)]
-struct CreatedFork<N: Network, SPEC, BLOCK = BlockEnv> {
+struct CreatedFork<N: Network, SPEC, BLOCK: ForkBlockEnv> {
     /// How the fork was initially created.
     opts: CreateFork,
     /// The resolved EVM environment (fetched from the provider).
     evm_env: EvmEnv<SPEC, BLOCK>,
     /// Copy of the sender.
-    backend: SharedBackend<N>,
+    backend: SharedBackend<N, BLOCK>,
     /// How many consumers there are, since a `SharedBacked` can be used by multiple
     /// consumers.
     num_senders: Arc<AtomicUsize>,
 }
 
-impl<N: Network, SPEC> CreatedFork<N, SPEC> {
-    pub fn new(opts: CreateFork, evm_env: EvmEnv<SPEC>, backend: SharedBackend<N>) -> Self {
+impl<N: Network, SPEC, BLOCK: ForkBlockEnv> CreatedFork<N, SPEC, BLOCK> {
+    pub fn new(
+        opts: CreateFork,
+        evm_env: EvmEnv<SPEC, BLOCK>,
+        backend: SharedBackend<N, BLOCK>,
+    ) -> Self {
         Self { opts, evm_env, backend, num_senders: Arc::new(AtomicUsize::new(1)) }
     }
 
@@ -531,11 +559,11 @@ impl<N: Network, SPEC> CreatedFork<N, SPEC> {
 /// This type intentionally does not implement `Clone` since it's intended that there's only once
 /// instance.
 #[derive(Debug)]
-struct ShutDownMultiFork<N: Network, SPEC, BLOCK = BlockEnv> {
+struct ShutDownMultiFork<N: Network, SPEC, BLOCK: ForkBlockEnv> {
     handler: Option<Sender<Request<N, SPEC, BLOCK>>>,
 }
 
-impl<N: Network, SPEC, BLOCK> Drop for ShutDownMultiFork<N, SPEC, BLOCK> {
+impl<N: Network, SPEC, BLOCK: ForkBlockEnv> Drop for ShutDownMultiFork<N, SPEC, BLOCK> {
     fn drop(&mut self) {
         trace!(target: "fork::multi", "initiating shutdown");
         let (sender, rx) = oneshot_channel();
@@ -552,17 +580,21 @@ impl<N: Network, SPEC, BLOCK> Drop for ShutDownMultiFork<N, SPEC, BLOCK> {
 /// Creates a new fork.
 ///
 /// This will establish a new `Provider` to the endpoint and return the Fork Backend.
-async fn create_fork<N: Network, SPEC: Into<SpecId> + Default + Copy>(
+async fn create_fork<
+    N: Network,
+    SPEC: Into<SpecId> + Default + Copy,
+    BLOCK: FoundryBlock + ForkBlockEnv + Default,
+>(
     mut fork: CreateFork,
-) -> eyre::Result<(ForkId, CreatedFork<N, SPEC>, BackendHandler<N>)> {
+) -> eyre::Result<(ForkId, CreatedFork<N, SPEC, BLOCK>, BackendHandler<N, BLOCK>)> {
     // Ensure evm_opts reflects the fork URL (may differ from the resolved CreateFork url when
     // created via cheatcodes, where evm_opts is cloned from the base config).
     fork.evm_opts.fork_url = Some(fork.url.clone());
 
     let provider = fork.evm_opts.fork_provider_with_url::<N>(&fork.url)?;
 
     // Initialise the fork environment.
-    let (evm_env, number) = fork.evm_opts.fork_evm_env::<_, BlockEnv, _, _>(&provider).await?;
+    let (evm_env, number) = fork.evm_opts.fork_evm_env::<_, BLOCK, _, _>(&provider).await?;
     let meta = BlockchainDbMeta::new(evm_env.block_env.clone(), fork.url.clone());
 
     // Determine the cache path if caching is enabled.
