feat(executor): call hook with state changes after post block balance increments (#13050)

crates/ethereum/evm/src/execute.rs

@@ -13,8 +13,9 @@ use reth_consensus::ConsensusError;
 use reth_ethereum_consensus::validate_block_post_execution;
 use reth_evm::{
     execute::{
-        BasicBlockExecutorProvider, BlockExecutionError, BlockExecutionStrategy,
-        BlockExecutionStrategyFactory, BlockValidationError, ExecuteOutput, ProviderError,
+        balance_increment_state, BasicBlockExecutorProvider, BlockExecutionError,
+        BlockExecutionStrategy, BlockExecutionStrategyFactory, BlockValidationError, ExecuteOutput,
+        ProviderError,
     },
     state_change::post_block_balance_increments,
     system_calls::{OnStateHook, SystemCaller},
@@ -263,8 +264,11 @@ where
         }
         // increment balances
         self.state
-            .increment_balances(balance_increments)
+            .increment_balances(balance_increments.clone())
             .map_err(|_| BlockValidationError::IncrementBalanceFailed)?;
+        // call state hook with changes due to balance increments.
+        let balance_state = balance_increment_state(&balance_increments, &mut self.state)?;
+        self.system_caller.on_state(&balance_state);
 
         Ok(requests)
     }
@@ -317,6 +321,7 @@ mod tests {
     use alloy_eips::{
         eip2935::{HISTORY_STORAGE_ADDRESS, HISTORY_STORAGE_CODE},
         eip4788::{BEACON_ROOTS_ADDRESS, BEACON_ROOTS_CODE, SYSTEM_ADDRESS},
+        eip4895::Withdrawal,
         eip7002::{WITHDRAWAL_REQUEST_PREDEPLOY_ADDRESS, WITHDRAWAL_REQUEST_PREDEPLOY_CODE},
         eip7685::EMPTY_REQUESTS_HASH,
     };
@@ -333,9 +338,9 @@ mod tests {
         database::StateProviderDatabase, test_utils::StateProviderTest, TransitionState,
     };
     use reth_testing_utils::generators::{self, sign_tx_with_key_pair};
-    use revm_primitives::BLOCKHASH_SERVE_WINDOW;
+    use revm_primitives::{address, EvmState, BLOCKHASH_SERVE_WINDOW};
     use secp256k1::{Keypair, Secp256k1};
-    use std::collections::HashMap;
+    use std::{collections::HashMap, sync::mpsc};
 
     fn create_state_provider_with_beacon_root_contract() -> StateProviderTest {
         let mut db = StateProviderTest::default();
@@ -1220,4 +1225,67 @@ mod tests {
             ),
         }
     }
+
+    #[test]
+    fn test_balance_increment_not_duplicated() {
+        let chain_spec = Arc::new(
+            ChainSpecBuilder::from(&*MAINNET)
+                .shanghai_activated()
+                .with_fork(EthereumHardfork::Prague, ForkCondition::Timestamp(0))
+                .build(),
+        );
+
+        let withdrawal_recipient = address!("1000000000000000000000000000000000000000");
+
+        let mut db = StateProviderTest::default();
+        let initial_balance = 100;
+        db.insert_account(
+            withdrawal_recipient,
+            Account { balance: U256::from(initial_balance), nonce: 1, bytecode_hash: None },
+            None,
+            HashMap::default(),
+        );
+
+        let withdrawal =
+            Withdrawal { index: 0, validator_index: 0, address: withdrawal_recipient, amount: 1 };
+
+        let header = Header { timestamp: 1, number: 1, ..Header::default() };
+
+        let block = BlockWithSenders {
+            block: Block {
+                header,
+                body: BlockBody {
+                    transactions: vec![],
+                    ommers: vec![],
+                    withdrawals: Some(vec![withdrawal].into()),
+                },
+            },
+            senders: vec![],
+        };
+
+        let provider = executor_provider(chain_spec);
+        let executor = provider.executor(StateProviderDatabase::new(&db));
+
+        let (tx, rx) = mpsc::channel();
+        let tx_clone = tx.clone();
+
+        let _output = executor
+            .execute_with_state_hook((&block, U256::ZERO).into(), move |state: &EvmState| {
+                if let Some(account) = state.get(&withdrawal_recipient) {
+                    let _ = tx_clone.send(account.info.balance);
+                }
+            })
+            .expect("Block execution should succeed");
+
+        drop(tx);
+        let balance_changes: Vec<U256> = rx.try_iter().collect();
+
+        if let Some(final_balance) = balance_changes.last() {
+            let expected_final_balance = U256::from(initial_balance) + U256::from(1_000_000_000); // initial + 1 Gwei in Wei
+            assert_eq!(
+                *final_balance, expected_final_balance,
+                "Final balance should match expected value after withdrawal"
+            );
+        }
+    }
 }

crates/evm/src/execute.rs

@@ -12,7 +12,10 @@ pub use reth_storage_errors::provider::ProviderError;
 use crate::{system_calls::OnStateHook, TxEnvOverrides};
 use alloc::{boxed::Box, vec::Vec};
 use alloy_eips::eip7685::Requests;
-use alloy_primitives::BlockNumber;
+use alloy_primitives::{
+    map::{DefaultHashBuilder, HashMap},
+    Address, BlockNumber,
+};
 use core::fmt::Display;
 use reth_consensus::ConsensusError;
 use reth_primitives::{BlockWithSenders, NodePrimitives, Receipt};
@@ -22,7 +25,7 @@ use revm::{
     db::{states::bundle_state::BundleRetention, BundleState},
     State,
 };
-use revm_primitives::{db::Database, U256};
+use revm_primitives::{db::Database, Account, AccountStatus, EvmState, U256};
 
 /// A general purpose executor trait that executes an input (e.g. block) and produces an output
 /// (e.g. state changes and receipts).
@@ -499,6 +502,42 @@ where
     }
 }
 
+/// Creates an `EvmState` from a map of balance increments and the current state
+/// to load accounts from. No balance increment is done in the function.
+/// Zero balance increments are ignored and won't create state entries.
+pub fn balance_increment_state<DB>(
+    balance_increments: &HashMap<Address, u128, DefaultHashBuilder>,
+    state: &mut State<DB>,
+) -> Result<EvmState, BlockExecutionError>
+where
+    DB: Database,
+{
+    let mut load_account = |address: &Address| -> Result<(Address, Account), BlockExecutionError> {
+        let cache_account = state.load_cache_account(*address).map_err(|_| {
+            BlockExecutionError::msg("could not load account for balance increment")
+        })?;
+
+        let account = cache_account.account.as_ref().ok_or_else(|| {
+            BlockExecutionError::msg("could not load account for balance increment")
+        })?;
+
+        Ok((
+            *address,
+            Account {
+                info: account.info.clone(),
+                storage: Default::default(),
+                status: AccountStatus::Touched,
+            },
+        ))
+    };
+
+    balance_increments
+        .iter()
+        .filter(|(_, &balance)| balance != 0)
+        .map(|(addr, _)| load_account(addr))
+        .collect::<Result<EvmState, _>>()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -507,7 +546,7 @@ mod tests {
     use reth_chainspec::{ChainSpec, MAINNET};
     use reth_primitives::EthPrimitives;
     use revm::db::{CacheDB, EmptyDBTyped};
-    use revm_primitives::{bytes, TxEnv};
+    use revm_primitives::{address, bytes, AccountInfo, TxEnv, KECCAK_EMPTY};
     use std::sync::Arc;
 
     #[derive(Clone, Default)]
@@ -760,4 +799,90 @@ mod tests {
         let result = executor.execute(BlockExecutionInput::new(&Default::default(), U256::ZERO));
         assert!(result.is_ok());
     }
+
+    fn setup_state_with_account(
+        addr: Address,
+        balance: u128,
+        nonce: u64,
+    ) -> State<CacheDB<EmptyDBTyped<BlockExecutionError>>> {
+        let db = CacheDB::<EmptyDBTyped<BlockExecutionError>>::default();
+        let mut state = State::builder().with_database(db).with_bundle_update().build();
+
+        let account_info = AccountInfo {
+            balance: U256::from(balance),
+            nonce,
+            code_hash: KECCAK_EMPTY,
+            code: None,
+        };
+        state.insert_account(addr, account_info);
+        state
+    }
+
+    #[test]
+    fn test_balance_increment_state_zero() {
+        let addr = address!("1000000000000000000000000000000000000000");
+        let mut state = setup_state_with_account(addr, 100, 1);
+
+        let mut increments = HashMap::<Address, u128, DefaultHashBuilder>::default();
+        increments.insert(addr, 0);
+
+        let result = balance_increment_state(&increments, &mut state).unwrap();
+        assert!(result.is_empty(), "Zero increments should be ignored");
+    }
+
+    #[test]
+    fn test_balance_increment_state_empty_increments_map() {
+        let mut state = State::builder()
+            .with_database(CacheDB::<EmptyDBTyped<BlockExecutionError>>::default())
+            .with_bundle_update()
+            .build();
+
+        let increments = HashMap::<Address, u128, DefaultHashBuilder>::default();
+        let result = balance_increment_state(&increments, &mut state).unwrap();
+        assert!(result.is_empty(), "Empty increments map should return empty state");
+    }
+
+    #[test]
+    fn test_balance_increment_state_multiple_valid_increments() {
+        let addr1 = address!("1000000000000000000000000000000000000000");
+        let addr2 = address!("2000000000000000000000000000000000000000");
+
+        let mut state = setup_state_with_account(addr1, 100, 1);
+
+        let account2 =
+            AccountInfo { balance: U256::from(200), nonce: 1, code_hash: KECCAK_EMPTY, code: None };
+        state.insert_account(addr2, account2);
+
+        let mut increments = HashMap::<Address, u128, DefaultHashBuilder>::default();
+        increments.insert(addr1, 50);
+        increments.insert(addr2, 100);
+
+        let result = balance_increment_state(&increments, &mut state).unwrap();
+
+        assert_eq!(result.len(), 2);
+        assert_eq!(result.get(&addr1).unwrap().info.balance, U256::from(100));
+        assert_eq!(result.get(&addr2).unwrap().info.balance, U256::from(200));
+    }
+
+    #[test]
+    fn test_balance_increment_state_mixed_zero_and_nonzero_increments() {
+        let addr1 = address!("1000000000000000000000000000000000000000");
+        let addr2 = address!("2000000000000000000000000000000000000000");
+
+        let mut state = setup_state_with_account(addr1, 100, 1);
+
+        let account2 =
+            AccountInfo { balance: U256::from(200), nonce: 1, code_hash: KECCAK_EMPTY, code: None };
+        state.insert_account(addr2, account2);
+
+        let mut increments = HashMap::<Address, u128, DefaultHashBuilder>::default();
+        increments.insert(addr1, 0);
+        increments.insert(addr2, 100);
+
+        let result = balance_increment_state(&increments, &mut state).unwrap();
+
+        assert_eq!(result.len(), 1, "Only non-zero increments should be included");
+        assert!(!result.contains_key(&addr1), "Zero increment account should not be included");
+        assert_eq!(result.get(&addr2).unwrap().info.balance, U256::from(200));
+    }
 }

crates/optimism/evm/src/execute.rs

@@ -10,8 +10,9 @@ use reth_chainspec::EthereumHardforks;
 use reth_consensus::ConsensusError;
 use reth_evm::{
     execute::{
-        BasicBlockExecutorProvider, BlockExecutionError, BlockExecutionStrategy,
-        BlockExecutionStrategyFactory, BlockValidationError, ExecuteOutput, ProviderError,
+        balance_increment_state, BasicBlockExecutorProvider, BlockExecutionError,
+        BlockExecutionStrategy, BlockExecutionStrategyFactory, BlockValidationError, ExecuteOutput,
+        ProviderError,
     },
     state_change::post_block_balance_increments,
     system_calls::{OnStateHook, SystemCaller},
@@ -260,8 +261,11 @@ where
             post_block_balance_increments(&self.chain_spec.clone(), block, total_difficulty);
         // increment balances
         self.state
-            .increment_balances(balance_increments)
+            .increment_balances(balance_increments.clone())
             .map_err(|_| BlockValidationError::IncrementBalanceFailed)?;
+        // call state hook with changes due to balance increments.
+        let balance_state = balance_increment_state(&balance_increments, &mut self.state)?;
+        self.system_caller.on_state(&balance_state);
 
         Ok(Requests::default())
     }