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"
+            );
+        }
+    }
 }