fix(tree): reinsert unwound state to dependent chains (#5683)

Co-authored-by: Roman Krasiuk <rokrassyuk@gmail.com>

Cargo.lock

@@ -5691,6 +5691,7 @@ dependencies = [
  "reth-metrics",
  "reth-primitives",
  "reth-provider",
+ "reth-revm",
  "reth-stages",
  "tokio",
  "tracing",

bin/reth/Cargo.toml

@@ -131,6 +131,7 @@ optimism = [
     "reth-basic-payload-builder/optimism",
     "reth-network/optimism",
     "reth-network-api/optimism",
+    "reth-blockchain-tree/optimism",
 ]
 # no-op feature flag for switching between the `optimism` and default functionality in CI matrices
 ethereum = []

crates/blockchain-tree/Cargo.toml

@@ -40,8 +40,10 @@ reth-db = { workspace = true, features = ["test-utils"] }
 reth-interfaces = { workspace = true, features = ["test-utils"] }
 reth-primitives = { workspace = true , features = ["test-utils"] }
 reth-provider = { workspace = true, features = ["test-utils"] }
+reth-revm.workspace = true
 parking_lot.workspace = true
 assert_matches.workspace = true
 
 [features]
 test-utils = []
+optimism = ["reth-primitives/optimism", "reth-interfaces/optimism", "reth-provider/optimism", "reth-revm/optimism"]

crates/blockchain-tree/src/blockchain_tree.rs

@@ -28,7 +28,10 @@ use reth_provider::{
     ChainSpecProvider, DisplayBlocksChain, ExecutorFactory, HeaderProvider, ProviderError,
 };
 use reth_stages::{MetricEvent, MetricEventsSender};
-use std::{collections::BTreeMap, sync::Arc};
+use std::{
+    collections::{BTreeMap, HashSet},
+    sync::Arc,
+};
 use tracing::{debug, error, info, instrument, trace, warn};
 
 #[cfg_attr(doc, aquamarine::aquamarine)]
@@ -607,6 +610,70 @@ impl<DB: Database, EF: ExecutorFactory> BlockchainTree<DB, EF> {
         self.state.insert_chain(chain)
     }
 
+    /// Iterate over all child chains that depend on this block and return
+    /// their ids.
+    fn find_all_dependent_chains(&self, block: &BlockHash) -> HashSet<BlockChainId> {
+        // Find all forks of given block.
+        let mut dependent_block =
+            self.block_indices().fork_to_child().get(block).cloned().unwrap_or_default();
+        let mut dependent_chains = HashSet::new();
+
+        while let Some(block) = dependent_block.pop_back() {
+            // Get chain of dependent block.
+            let chain_id =
+                self.block_indices().get_blocks_chain_id(&block).expect("Block should be in tree");
+
+            // Find all blocks that fork from this chain.
+            for chain_block in
+                self.state.chains.get(&chain_id).expect("Chain should be in tree").blocks().values()
+            {
+                if let Some(forks) = self.block_indices().fork_to_child().get(&chain_block.hash()) {
+                    // If there are sub forks append them for processing.
+                    dependent_block.extend(forks);
+                }
+            }
+            // Insert dependent chain id.
+            dependent_chains.insert(chain_id);
+        }
+        dependent_chains
+    }
+
+    /// Inserts unwound chain back into the tree and updates any dependent chains.
+    ///
+    /// This method searches for any chain that depended on this block being part of the canonical
+    /// chain. Each dependent chain's state is then updated with state entries removed from the
+    /// plain state during the unwind.
+    fn insert_unwound_chain(&mut self, chain: AppendableChain) -> Option<BlockChainId> {
+        // iterate over all blocks in chain and find any fork blocks that are in tree.
+        for (number, block) in chain.blocks().iter() {
+            let hash = block.hash();
+
+            // find all chains that fork from this block.
+            let chains_to_bump = self.find_all_dependent_chains(&hash);
+            if !chains_to_bump.is_empty() {
+                // if there is such chain, revert state to this block.
+                let mut cloned_state = chain.state().clone();
+                cloned_state.revert_to(*number);
+
+                // prepend state to all chains that fork from this block.
+                for chain_id in chains_to_bump {
+                    let chain =
+                        self.state.chains.get_mut(&chain_id).expect("Chain should be in tree");
+
+                    debug!(target: "blockchain_tree",
+                        unwound_block= ?block.num_hash(),
+                        chain_id = ?chain_id,
+                        chain_tip = ?chain.tip().num_hash(),
+                        "Prepend unwound block state to blockchain tree chain");
+
+                    chain.prepend_state(cloned_state.state().clone())
+                }
+            }
+        }
+        // Insert unwound chain to the tree.
+        self.insert_chain(chain)
+    }
+
     /// Checks the block buffer for the given block.
     pub fn get_buffered_block(&self, hash: &BlockHash) -> Option<&SealedBlockWithSenders> {
         self.state.get_buffered_block(hash)
@@ -1064,7 +1131,7 @@ impl<DB: Database, EF: ExecutorFactory> BlockchainTree<DB, EF> {
                 let reorg_depth = old_canon_chain.len();
 
                 // insert old canon chain
-                self.insert_chain(AppendableChain::new(old_canon_chain));
+                self.insert_unwound_chain(AppendableChain::new(old_canon_chain));
                 durations_recorder.record_relative(MakeCanonicalAction::InsertOldCanonicalChain);
 
                 self.update_reorg_metrics(reorg_depth as f64);
@@ -1156,7 +1223,7 @@ impl<DB: Database, EF: ExecutorFactory> BlockchainTree<DB, EF> {
         if let Some(old_canon_chain) = old_canon_chain {
             self.block_indices_mut().unwind_canonical_chain(unwind_to);
             // insert old canonical chain to BlockchainTree.
-            self.insert_chain(AppendableChain::new(old_canon_chain));
+            self.insert_unwound_chain(AppendableChain::new(old_canon_chain));
         }
 
         Ok(())
@@ -1229,7 +1296,14 @@ mod tests {
     use reth_db::{tables, test_utils::TempDatabase, transaction::DbTxMut, DatabaseEnv};
     use reth_interfaces::test_utils::TestConsensus;
     use reth_primitives::{
-        constants::EMPTY_ROOT_HASH, stage::StageCheckpoint, ChainSpecBuilder, B256, MAINNET,
+        constants::{EIP1559_INITIAL_BASE_FEE, EMPTY_ROOT_HASH, ETHEREUM_BLOCK_GAS_LIMIT},
+        keccak256,
+        proofs::{calculate_receipt_root, calculate_transaction_root, state_root_unhashed},
+        revm_primitives::AccountInfo,
+        stage::StageCheckpoint,
+        Account, Address, ChainSpecBuilder, Genesis, GenesisAccount, Header, Signature,
+        Transaction, TransactionKind, TransactionSigned, TransactionSignedEcRecovered, TxEip1559,
+        B256, MAINNET,
     };
     use reth_provider::{
         test_utils::{
@@ -1238,6 +1312,7 @@ mod tests {
         },
         BlockWriter, BundleStateWithReceipts, ProviderFactory,
     };
+    use reth_revm::EvmProcessorFactory;
     use std::{
         collections::{HashMap, HashSet},
         sync::Arc,
@@ -1354,6 +1429,196 @@ mod tests {
         }
     }
 
+    #[test]
+    fn consecutive_reorgs() {
+        let signer = Address::random();
+        let initial_signer_balance = U256::from(10).pow(U256::from(18));
+        let chain_spec = Arc::new(
+            ChainSpecBuilder::default()
+                .chain(MAINNET.chain)
+                .genesis(Genesis {
+                    alloc: HashMap::from([(
+                        signer,
+                        GenesisAccount { balance: initial_signer_balance, ..Default::default() },
+                    )]),
+                    ..MAINNET.genesis.clone()
+                })
+                .shanghai_activated()
+                .build(),
+        );
+        let provider_factory = create_test_provider_factory_with_chain_spec(chain_spec.clone());
+        let consensus = Arc::new(TestConsensus::default());
+        let executor_factory = EvmProcessorFactory::new(chain_spec.clone());
+
+        {
+            let provider_rw = provider_factory.provider_rw().unwrap();
+            provider_rw
+                .insert_block(
+                    SealedBlock::new(chain_spec.sealed_genesis_header(), Default::default()),
+                    Some(Vec::new()),
+                    None,
+                )
+                .unwrap();
+            let account = Account { balance: initial_signer_balance, ..Default::default() };
+            provider_rw.tx_ref().put::<tables::PlainAccountState>(signer, account).unwrap();
+            provider_rw.tx_ref().put::<tables::HashedAccount>(keccak256(signer), account).unwrap();
+            provider_rw.commit().unwrap();
+        }
+
+        let single_tx_cost = U256::from(EIP1559_INITIAL_BASE_FEE * 21_000);
+        let mock_tx = |nonce: u64| -> TransactionSignedEcRecovered {
+            TransactionSigned::from_transaction_and_signature(
+                Transaction::Eip1559(TxEip1559 {
+                    chain_id: chain_spec.chain.id(),
+                    nonce,
+                    gas_limit: 21_000,
+                    to: TransactionKind::Call(Address::ZERO),
+                    max_fee_per_gas: EIP1559_INITIAL_BASE_FEE as u128,
+                    ..Default::default()
+                }),
+                Signature::default(),
+            )
+            .with_signer(signer)
+        };
+
+        let mock_block = |number: u64,
+                          parent: Option<B256>,
+                          body: Vec<TransactionSignedEcRecovered>,
+                          num_of_signer_txs: u64|
+         -> SealedBlockWithSenders {
+            let transactions_root = calculate_transaction_root(&body);
+            let receipts = body
+                .iter()
+                .enumerate()
+                .map(|(idx, tx)| {
+                    Receipt {
+                        tx_type: tx.tx_type(),
+                        success: true,
+                        cumulative_gas_used: (idx as u64 + 1) * 21_000,
+                        ..Default::default()
+                    }
+                    .with_bloom()
+                })
+                .collect::<Vec<_>>();
+
+            #[cfg(not(feature = "optimism"))]
+            let receipts_root = calculate_receipt_root(&receipts);
+
+            #[cfg(feature = "optimism")]
+            let receipts_root = calculate_receipt_root(&receipts, &chain_spec, 0);
+
+            SealedBlockWithSenders::new(
+                SealedBlock {
+                    header: Header {
+                        number,
+                        parent_hash: parent.unwrap_or_default(),
+                        gas_used: body.len() as u64 * 21_000,
+                        gas_limit: ETHEREUM_BLOCK_GAS_LIMIT,
+                        mix_hash: B256::random(),
+                        base_fee_per_gas: Some(EIP1559_INITIAL_BASE_FEE),
+                        transactions_root,
+                        receipts_root,
+                        state_root: state_root_unhashed(HashMap::from([(
+                            signer,
+                            (
+                                AccountInfo {
+                                    balance: initial_signer_balance -
+                                        (single_tx_cost * U256::from(num_of_signer_txs)),
+                                    nonce: num_of_signer_txs,
+                                    ..Default::default()
+                                },
+                                EMPTY_ROOT_HASH,
+                            ),
+                        )])),
+                        ..Default::default()
+                    }
+                    .seal_slow(),
+                    body: body.clone().into_iter().map(|tx| tx.into_signed()).collect(),
+                    ommers: Vec::new(),
+                    withdrawals: Some(Vec::new()),
+                },
+                body.iter().map(|tx| tx.signer()).collect(),
+            )
+            .unwrap()
+        };
+
+        let fork_block = mock_block(1, Some(chain_spec.genesis_hash()), Vec::from([mock_tx(0)]), 1);
+
+        let canonical_block_1 =
+            mock_block(2, Some(fork_block.hash), Vec::from([mock_tx(1), mock_tx(2)]), 3);
+        let canonical_block_2 = mock_block(3, Some(canonical_block_1.hash), Vec::new(), 3);
+        let canonical_block_3 =
+            mock_block(4, Some(canonical_block_2.hash), Vec::from([mock_tx(3)]), 4);
+
+        let sidechain_block_1 = mock_block(2, Some(fork_block.hash), Vec::from([mock_tx(1)]), 2);
+        let sidechain_block_2 =
+            mock_block(3, Some(sidechain_block_1.hash), Vec::from([mock_tx(2)]), 3);
+
+        let mut tree = BlockchainTree::new(
+            TreeExternals::new(provider_factory.clone(), consensus, executor_factory.clone()),
+            BlockchainTreeConfig::default(),
+            None,
+        )
+        .expect("failed to create tree");
+
+        tree.insert_block(fork_block.clone(), BlockValidationKind::Exhaustive).unwrap();
+
+        assert_eq!(
+            tree.make_canonical(&fork_block.hash).unwrap(),
+            CanonicalOutcome::Committed { head: fork_block.header.clone() }
+        );
+
+        assert_eq!(
+            tree.insert_block(canonical_block_1.clone(), BlockValidationKind::Exhaustive).unwrap(),
+            InsertPayloadOk::Inserted(BlockStatus::Valid)
+        );
+
+        assert_eq!(
+            tree.make_canonical(&canonical_block_1.hash).unwrap(),
+            CanonicalOutcome::Committed { head: canonical_block_1.header.clone() }
+        );
+
+        assert_eq!(
+            tree.insert_block(canonical_block_2.clone(), BlockValidationKind::Exhaustive).unwrap(),
+            InsertPayloadOk::Inserted(BlockStatus::Valid)
+        );
+
+        assert_eq!(
+            tree.insert_block(sidechain_block_1.clone(), BlockValidationKind::Exhaustive).unwrap(),
+            InsertPayloadOk::Inserted(BlockStatus::Accepted)
+        );
+
+        assert_eq!(
+            tree.make_canonical(&sidechain_block_1.hash).unwrap(),
+            CanonicalOutcome::Committed { head: sidechain_block_1.header.clone() }
+        );
+
+        assert_eq!(
+            tree.make_canonical(&canonical_block_1.hash).unwrap(),
+            CanonicalOutcome::Committed { head: canonical_block_1.header.clone() }
+        );
+
+        assert_eq!(
+            tree.insert_block(sidechain_block_2.clone(), BlockValidationKind::Exhaustive).unwrap(),
+            InsertPayloadOk::Inserted(BlockStatus::Accepted)
+        );
+
+        assert_eq!(
+            tree.make_canonical(&sidechain_block_2.hash).unwrap(),
+            CanonicalOutcome::Committed { head: sidechain_block_2.header.clone() }
+        );
+
+        assert_eq!(
+            tree.insert_block(canonical_block_3.clone(), BlockValidationKind::Exhaustive).unwrap(),
+            InsertPayloadOk::Inserted(BlockStatus::Accepted)
+        );
+
+        assert_eq!(
+            tree.make_canonical(&canonical_block_3.hash).unwrap(),
+            CanonicalOutcome::Committed { head: canonical_block_3.header.clone() }
+        );
+    }
+
     #[tokio::test]
     async fn sanity_path() {
         let data = BlockChainTestData::default_with_numbers(11, 12);

crates/consensus/beacon/Cargo.toml

@@ -60,4 +60,5 @@ optimism = [
   "reth-rpc-types/optimism",
   "reth-rpc-types-compat/optimism",
   "reth-payload-builder/optimism",
+  "reth-blockchain-tree/optimism",
 ]

crates/primitives/src/trie/account.rs

@@ -3,6 +3,7 @@ use crate::{
 };
 use alloy_primitives::keccak256;
 use alloy_rlp::{RlpDecodable, RlpEncodable};
+use revm_primitives::AccountInfo;
 
 /// An Ethereum account as represented in the trie.
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Default, RlpEncodable, RlpDecodable)]
@@ -28,6 +29,17 @@ impl From<(Account, B256)> for TrieAccount {
     }
 }
 
+impl From<(AccountInfo, B256)> for TrieAccount {
+    fn from((account, storage_root): (AccountInfo, B256)) -> Self {
+        Self {
+            nonce: account.nonce,
+            balance: account.balance,
+            storage_root,
+            code_hash: account.code_hash,
+        }
+    }
+}
+
 impl From<GenesisAccount> for TrieAccount {
     fn from(account: GenesisAccount) -> Self {
         let storage_root = account

crates/storage/provider/src/bundle_state/bundle_state_with_receipts.rs

@@ -351,6 +351,22 @@ impl BundleStateWithReceipts {
         self.receipts.extend(other.receipts.receipt_vec);
     }
 
+    /// Prepends present the state with the given BundleState.
+    /// It adds changes from the given state but does not override any existing changes.
+    ///
+    /// Reverts  and receipts are not updated.
+    pub fn prepend_state(&mut self, mut other: BundleState) {
+        let other_len = other.reverts.len();
+        // take this bundle
+        let this_bundle = std::mem::take(&mut self.bundle);
+        // extend other bundle with this
+        other.extend(this_bundle);
+        // discard other reverts
+        other.take_n_reverts(other_len);
+        // swap bundles
+        std::mem::swap(&mut self.bundle, &mut other)
+    }
+
     /// Write bundle state to database.
     ///
     /// `omit_changed_check` should be set to true of bundle has some of it data
@@ -1355,4 +1371,42 @@ mod tests {
         state.merge_transitions(BundleRetention::PlainState);
         assert_state_root(&state, &prestate, "recreated changed storage");
     }
+
+    #[test]
+    fn prepend_state() {
+        let address1 = Address::random();
+        let address2 = Address::random();
+
+        let account1 = RevmAccountInfo { nonce: 1, ..Default::default() };
+        let account1_changed = RevmAccountInfo { nonce: 1, ..Default::default() };
+        let account2 = RevmAccountInfo { nonce: 1, ..Default::default() };
+
+        let present_state = BundleState::builder(2..=2)
+            .state_present_account_info(address1, account1_changed.clone())
+            .build();
+        assert_eq!(present_state.reverts.len(), 1);
+        let previous_state = BundleState::builder(1..=1)
+            .state_present_account_info(address1, account1)
+            .state_present_account_info(address2, account2.clone())
+            .build();
+        assert_eq!(previous_state.reverts.len(), 1);
+
+        let mut test = BundleStateWithReceipts {
+            bundle: present_state,
+            receipts: Receipts::from_vec(vec![vec![Some(Receipt::default()); 2]; 1]),
+            first_block: 2,
+        };
+
+        test.prepend_state(previous_state);
+
+        assert_eq!(test.receipts.len(), 1);
+        let end_state = test.state();
+        assert_eq!(end_state.state.len(), 2);
+        // reverts num should stay the same.
+        assert_eq!(end_state.reverts.len(), 1);
+        // account1 is not overwritten.
+        assert_eq!(end_state.state.get(&address1).unwrap().info, Some(account1_changed));
+        // account2 got inserted
+        assert_eq!(end_state.state.get(&address2).unwrap().info, Some(account2));
+    }
 }

crates/storage/provider/src/chain.rs

@@ -6,6 +6,7 @@ use reth_primitives::{
     Address, BlockHash, BlockNumHash, BlockNumber, ForkBlock, Receipt, SealedBlock,
     SealedBlockWithSenders, SealedHeader, TransactionSigned, TransactionSignedEcRecovered, TxHash,
 };
+use revm::db::BundleState;
 use std::{borrow::Cow, collections::BTreeMap, fmt};
 
 /// A chain of blocks and their final state.
@@ -59,6 +60,11 @@ impl Chain {
         &self.state
     }
 
+    /// Prepends the given state to the current state.
+    pub fn prepend_state(&mut self, state: BundleState) {
+        self.state.prepend_state(state);
+    }
+
     /// Return true if chain is empty and has no blocks.
     pub fn is_empty(&self) -> bool {
         self.blocks.is_empty()
@@ -426,7 +432,7 @@ mod tests {
 
     #[test]
     fn chain_append() {
-        let block = SealedBlockWithSenders::default();
+        let block: SealedBlockWithSenders = SealedBlockWithSenders::default();
         let block1_hash = B256::new([0x01; 32]);
         let block2_hash = B256::new([0x02; 32]);
         let block3_hash = B256::new([0x03; 32]);