feat(engine): prefetch withdrawal addresses in pre-warming (#21966)

Co-authored-by: mattsse <matt@paradigm.xyz>
Co-authored-by: Amp <amp@ampcode.com>
Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>

crates/engine/tree/src/tree/payload_processor/mod.rs

@@ -11,7 +11,7 @@ use crate::tree::{
     StateProviderBuilder, TreeConfig,
 };
 use alloy_eip7928::BlockAccessList;
-use alloy_eips::eip1898::BlockWithParent;
+use alloy_eips::{eip1898::BlockWithParent, eip4895::Withdrawal};
 use alloy_evm::block::StateChangeSource;
 use alloy_primitives::B256;
 use crossbeam_channel::{Receiver as CrossbeamReceiver, Sender as CrossbeamSender};
@@ -976,6 +976,9 @@ pub struct ExecutionEnv<Evm: ConfigureEvm> {
     /// Used to determine parallel worker count for prewarming.
     /// A value of 0 indicates the count is unknown.
     pub transaction_count: usize,
+    /// Withdrawals included in the block.
+    /// Used to generate prefetch targets for withdrawal addresses.
+    pub withdrawals: Option<Vec<Withdrawal>>,
 }
 
 impl<Evm: ConfigureEvm> Default for ExecutionEnv<Evm>
@@ -989,6 +992,7 @@ where
             parent_hash: Default::default(),
             parent_state_root: Default::default(),
             transaction_count: 0,
+            withdrawals: None,
         }
     }
 }

crates/engine/tree/src/tree/payload_processor/prewarm.rs

@@ -24,6 +24,7 @@ use crate::tree::{
 };
 use alloy_consensus::transaction::TxHashRef;
 use alloy_eip7928::BlockAccessList;
+use alloy_eips::eip4895::Withdrawal;
 use alloy_evm::Database;
 use alloy_primitives::{keccak256, map::B256Set, B256};
 use crossbeam_channel::{Receiver as CrossbeamReceiver, Sender as CrossbeamSender};
@@ -163,7 +164,7 @@ where
             };
 
             // Spawn workers
-            let tx_sender = ctx.clone().spawn_workers(workers_needed, &executor, to_multi_proof, done_tx.clone());
+            let tx_sender = ctx.clone().spawn_workers(workers_needed, &executor,  to_multi_proof.clone(), done_tx.clone());
 
             // Distribute transactions to workers
             let mut tx_index = 0usize;
@@ -187,6 +188,16 @@ where
                 tx_index += 1;
             }
 
+            // Send withdrawal prefetch targets after all transactions have been distributed
+            if let Some(to_multi_proof) = to_multi_proof
+                && let Some(withdrawals) = &ctx.env.withdrawals
+                && !withdrawals.is_empty()
+            {
+                let targets =
+                    multiproof_targets_from_withdrawals(withdrawals, ctx.v2_proofs_enabled);
+                let _ = to_multi_proof.send(MultiProofMessage::PrefetchProofs(targets));
+            }
+
             // drop sender and wait for all tasks to finish
             drop(done_tx);
             drop(tx_sender);
@@ -831,6 +842,27 @@ fn multiproof_targets_v2_from_state(state: EvmState) -> (VersionedMultiProofTarg
     (VersionedMultiProofTargets::V2(targets), storage_target_count)
 }
 
+/// Returns [`VersionedMultiProofTargets`] for withdrawal addresses.
+///
+/// Withdrawals only modify account balances (no storage), so the targets contain
+/// only account-level entries with empty storage sets.
+fn multiproof_targets_from_withdrawals(
+    withdrawals: &[Withdrawal],
+    v2_enabled: bool,
+) -> VersionedMultiProofTargets {
+    use reth_trie_parallel::targets_v2::MultiProofTargetsV2;
+    if v2_enabled {
+        VersionedMultiProofTargets::V2(MultiProofTargetsV2 {
+            account_targets: withdrawals.iter().map(|w| keccak256(w.address).into()).collect(),
+            ..Default::default()
+        })
+    } else {
+        VersionedMultiProofTargets::Legacy(
+            withdrawals.iter().map(|w| (keccak256(w.address), Default::default())).collect(),
+        )
+    }
+}
+
 /// The events the pre-warm task can handle.
 ///
 /// Generic over `R` (receipt type) to allow sharing `Arc<ExecutionOutcome<R>>` with the main

crates/engine/tree/src/tree/payload_validator.rs

@@ -12,7 +12,7 @@ use crate::tree::{
 };
 use alloy_consensus::transaction::{Either, TxHashRef};
 use alloy_eip7928::BlockAccessList;
-use alloy_eips::{eip1898::BlockWithParent, NumHash};
+use alloy_eips::{eip1898::BlockWithParent, eip4895::Withdrawal, NumHash};
 use alloy_evm::Evm;
 use alloy_primitives::B256;
 
@@ -408,6 +408,7 @@ where
             parent_hash: input.parent_hash(),
             parent_state_root: parent_block.state_root(),
             transaction_count: input.transaction_count(),
+            withdrawals: input.withdrawals().map(|w| w.to_vec()),
         };
 
         // Plan the strategy used for state root computation.
@@ -1556,4 +1557,15 @@ impl<T: PayloadTypes> BlockOrPayload<T> {
             Self::Block(block) => block.transaction_count(),
         }
     }
+
+    /// Returns the withdrawals from the payload or block.
+    pub fn withdrawals(&self) -> Option<&[Withdrawal]>
+    where
+        T::ExecutionData: ExecutionPayload,
+    {
+        match self {
+            Self::Payload(payload) => payload.withdrawals().map(|w| w.as_slice()),
+            Self::Block(block) => block.body().withdrawals().map(|w| w.as_slice()),
+        }
+    }
 }