feat(tree): re-enable trie caching (#6276)

crates/blockchain-tree/src/blockchain_tree.rs

@@ -1258,6 +1258,7 @@ mod tests {
         BlockWriter, BundleStateWithReceipts, ProviderFactory,
     };
     use reth_revm::EvmProcessorFactory;
+    use reth_trie::StateRoot;
     use std::{
         collections::{HashMap, HashSet},
         sync::Arc,
@@ -1576,6 +1577,97 @@ mod tests {
         );
     }
 
+    #[test]
+    fn cached_trie_updates() {
+        let data = BlockChainTestData::default_from_number(11);
+        let (block1, exec1) = data.blocks[0].clone();
+        let (block2, exec2) = data.blocks[1].clone();
+        let (block3, exec3) = data.blocks[2].clone();
+        let (block4, exec4) = data.blocks[3].clone();
+        let (block5, exec5) = data.blocks[4].clone();
+        let genesis = data.genesis;
+
+        // test pops execution results from vector, so order is from last to first.
+        let externals = setup_externals(vec![exec5.clone(), exec4.clone(), exec3, exec2, exec1]);
+
+        // last finalized block would be number 9.
+        setup_genesis(&externals.provider_factory, genesis);
+
+        // make tree
+        let config = BlockchainTreeConfig::new(1, 2, 3, 2);
+        let mut tree = BlockchainTree::new(externals, config, None).expect("failed to create tree");
+        // genesis block 10 is already canonical
+        tree.make_canonical(&B256::ZERO).unwrap();
+
+        // make genesis block 10 as finalized
+        tree.finalize_block(10);
+
+        assert_eq!(
+            tree.insert_block(block1.clone(), BlockValidationKind::Exhaustive).unwrap(),
+            InsertPayloadOk::Inserted(BlockStatus::Valid(BlockAttachment::Canonical))
+        );
+        let block1_chain_id = tree.state.block_indices.get_blocks_chain_id(&block1.hash).unwrap();
+        let block1_chain = tree.state.chains.get(&block1_chain_id).unwrap();
+        assert!(block1_chain.trie_updates().is_some());
+
+        assert_eq!(
+            tree.insert_block(block2.clone(), BlockValidationKind::Exhaustive).unwrap(),
+            InsertPayloadOk::Inserted(BlockStatus::Valid(BlockAttachment::Canonical))
+        );
+        let block2_chain_id = tree.state.block_indices.get_blocks_chain_id(&block2.hash).unwrap();
+        let block2_chain = tree.state.chains.get(&block2_chain_id).unwrap();
+        assert!(block2_chain.trie_updates().is_some());
+
+        assert_eq!(
+            tree.make_canonical(&block2.hash).unwrap(),
+            CanonicalOutcome::Committed { head: block2.header.clone() }
+        );
+
+        assert_eq!(
+            tree.insert_block(block3.clone(), BlockValidationKind::Exhaustive).unwrap(),
+            InsertPayloadOk::Inserted(BlockStatus::Valid(BlockAttachment::Canonical))
+        );
+        let block3_chain_id = tree.state.block_indices.get_blocks_chain_id(&block3.hash).unwrap();
+        let block3_chain = tree.state.chains.get(&block3_chain_id).unwrap();
+        assert!(block3_chain.trie_updates().is_some());
+
+        assert_eq!(
+            tree.make_canonical(&block3.hash).unwrap(),
+            CanonicalOutcome::Committed { head: block3.header.clone() }
+        );
+
+        assert_eq!(
+            tree.insert_block(block4.clone(), BlockValidationKind::Exhaustive).unwrap(),
+            InsertPayloadOk::Inserted(BlockStatus::Valid(BlockAttachment::Canonical))
+        );
+        let block4_chain_id = tree.state.block_indices.get_blocks_chain_id(&block4.hash).unwrap();
+        let block4_chain = tree.state.chains.get(&block4_chain_id).unwrap();
+        assert!(block4_chain.trie_updates().is_some());
+
+        assert_eq!(
+            tree.insert_block(block5.clone(), BlockValidationKind::Exhaustive).unwrap(),
+            InsertPayloadOk::Inserted(BlockStatus::Valid(BlockAttachment::Canonical))
+        );
+
+        let block5_chain_id = tree.state.block_indices.get_blocks_chain_id(&block5.hash).unwrap();
+        let block5_chain = tree.state.chains.get(&block5_chain_id).unwrap();
+        assert!(block5_chain.trie_updates().is_some());
+
+        assert_eq!(
+            tree.make_canonical(&block5.hash).unwrap(),
+            CanonicalOutcome::Committed { head: block5.header.clone() }
+        );
+
+        let provider = tree.externals.provider_factory.provider().unwrap();
+        let (acc_prefix_set, storage_prefix_set) = exec5.hash_state_slow().construct_prefix_sets();
+        let state_root = StateRoot::from_tx(provider.tx_ref())
+            .with_changed_account_prefixes(acc_prefix_set)
+            .with_changed_storage_prefixes(storage_prefix_set)
+            .root()
+            .unwrap();
+        assert_eq!(state_root, block5.state_root);
+    }
+
     #[test]
     fn test_side_chain_fork() {
         let data = BlockChainTestData::default_from_number(11);

crates/blockchain-tree/src/chain.rs

@@ -167,7 +167,7 @@ impl AppendableChain {
         parent_block: &SealedHeader,
         bundle_state_data_provider: BSDP,
         externals: &TreeExternals<DB, EF>,
-        _block_attachment: BlockAttachment,
+        block_attachment: BlockAttachment,
         block_validation_kind: BlockValidationKind,
     ) -> RethResult<(BundleStateWithReceipts, Option<TrieUpdates>)>
     where
@@ -194,16 +194,13 @@ impl AppendableChain {
         // validation was requested.
         if block_validation_kind.is_exhaustive() {
             // check state root
-            // TODO: state root caching is disabled until debugged properly
-            // let (state_root, trie_updates) = if block_attachment.is_canonical() {
-            //     provider
-            //         .state_root_with_updates(&bundle_state)
-            //         .map(|(root, updates)| (root, Some(updates)))?
-
-            // } else {
-            //     (provider.state_root(&bundle_state)?, None)
-            // };
-            let state_root = provider.state_root(&bundle_state)?;
+            let (state_root, trie_updates) = if block_attachment.is_canonical() {
+                provider
+                    .state_root_with_updates(&bundle_state)
+                    .map(|(root, updates)| (root, Some(updates)))?
+            } else {
+                (provider.state_root(&bundle_state)?, None)
+            };
             if block.state_root != state_root {
                 return Err(ConsensusError::BodyStateRootDiff(
                     GotExpected { got: state_root, expected: block.state_root }.into(),
@@ -211,7 +208,7 @@ impl AppendableChain {
                 .into())
             }
 
-            Ok((bundle_state, None))
+            Ok((bundle_state, trie_updates))
         } else {
             Ok((bundle_state, None))
         }

crates/storage/provider/src/test_utils/blocks.rs

@@ -69,16 +69,30 @@ impl BlockChainTestData {
     /// Create test data with two blocks that are connected, specifying their block numbers.
     pub fn default_from_number(first: BlockNumber) -> Self {
         let one = block1(first);
-        let two = block2(first + 1, one.0.hash, &one.1);
-        Self { genesis: genesis(), blocks: vec![one, two] }
+        let mut extended_state = one.1.clone();
+        let two = block2(first + 1, one.0.hash, &extended_state);
+        extended_state.extend(two.1.clone());
+        let three = block3(first + 2, two.0.hash, &extended_state);
+        extended_state.extend(three.1.clone());
+        let four = block4(first + 3, three.0.hash, &extended_state);
+        extended_state.extend(four.1.clone());
+        let five = block5(first + 4, four.0.hash, &extended_state);
+        Self { genesis: genesis(), blocks: vec![one, two, three, four, five] }
     }
 }
 
 impl Default for BlockChainTestData {
     fn default() -> Self {
         let one = block1(1);
-        let two = block2(2, one.0.hash, &one.1);
-        Self { genesis: genesis(), blocks: vec![one, two] }
+        let mut extended_state = one.1.clone();
+        let two = block2(2, one.0.hash, &extended_state);
+        extended_state.extend(two.1.clone());
+        let three = block3(3, two.0.hash, &extended_state);
+        extended_state.extend(three.1.clone());
+        let four = block4(4, three.0.hash, &extended_state);
+        extended_state.extend(four.1.clone());
+        let five = block5(5, four.0.hash, &extended_state);
+        Self { genesis: genesis(), blocks: vec![one, two, three, four, five] }
     }
 }
 
@@ -104,6 +118,7 @@ fn bundle_state_root(state: &BundleStateWithReceipts) -> B256 {
                         account
                             .storage
                             .iter()
+                            .filter(|(_, value)| !value.present_value.is_zero())
                             .map(|(slot, value)| ((*slot).into(), value.present_value)),
                     ),
                 ),
@@ -222,3 +237,237 @@ fn block2(
 
     (SealedBlockWithSenders { block, senders: vec![Address::new([0x31; 20])] }, bundle)
 }
+
+/// Block three that points to block 2
+fn block3(
+    number: BlockNumber,
+    parent_hash: B256,
+    prev_state: &BundleStateWithReceipts,
+) -> (SealedBlockWithSenders, BundleStateWithReceipts) {
+    let address_range = 1..=20;
+    let slot_range = 1..=100;
+
+    let mut bundle_state_builder = BundleState::builder(number..=number);
+    for idx in address_range {
+        let address = Address::with_last_byte(idx);
+        bundle_state_builder = bundle_state_builder
+            .state_present_account_info(
+                address,
+                AccountInfo { nonce: 1, balance: U256::from(idx), ..Default::default() },
+            )
+            .state_storage(
+                address,
+                HashMap::from_iter(
+                    slot_range
+                        .clone()
+                        .map(|slot| (U256::from(slot), (U256::ZERO, U256::from(slot)))),
+                ),
+            );
+    }
+    let bundle = BundleStateWithReceipts::new(
+        bundle_state_builder.build(),
+        Receipts::from_vec(vec![vec![Some(Receipt {
+            tx_type: TxType::EIP1559,
+            success: true,
+            cumulative_gas_used: 400,
+            logs: vec![Log {
+                address: Address::new([0x61; 20]),
+                topics: vec![B256::with_last_byte(3), B256::with_last_byte(4)],
+                data: Bytes::default(),
+            }],
+            #[cfg(feature = "optimism")]
+            deposit_nonce: None,
+            #[cfg(feature = "optimism")]
+            deposit_receipt_version: None,
+        })]]),
+        number,
+    );
+
+    let mut extended = prev_state.clone();
+    extended.extend(bundle.clone());
+    let state_root = bundle_state_root(&extended);
+
+    let mut block = SealedBlock::decode(&mut BLOCK_RLP.as_slice()).unwrap();
+    block.withdrawals = Some(Withdrawals::new(vec![Withdrawal::default()]));
+    let mut header = block.header.clone().unseal();
+    header.number = number;
+    header.state_root = state_root;
+    // parent_hash points to block1 hash
+    header.parent_hash = parent_hash;
+    block.header = header.seal_slow();
+
+    (SealedBlockWithSenders { block, senders: vec![Address::new([0x31; 20])] }, bundle)
+}
+
+/// Block four that points to block 3
+fn block4(
+    number: BlockNumber,
+    parent_hash: B256,
+    prev_state: &BundleStateWithReceipts,
+) -> (SealedBlockWithSenders, BundleStateWithReceipts) {
+    let address_range = 1..=20;
+    let slot_range = 1..=100;
+
+    let mut bundle_state_builder = BundleState::builder(number..=number);
+    for idx in address_range {
+        let address = Address::with_last_byte(idx);
+        // increase balance for every even account and destroy every odd
+        bundle_state_builder = if idx % 2 == 0 {
+            bundle_state_builder
+                .state_present_account_info(
+                    address,
+                    AccountInfo { nonce: 1, balance: U256::from(idx * 2), ..Default::default() },
+                )
+                .state_storage(
+                    address,
+                    HashMap::from_iter(
+                        slot_range.clone().map(|slot| {
+                            (U256::from(slot), (U256::from(slot), U256::from(slot * 2)))
+                        }),
+                    ),
+                )
+        } else {
+            bundle_state_builder.state_address(address).state_storage(
+                address,
+                HashMap::from_iter(
+                    slot_range
+                        .clone()
+                        .map(|slot| (U256::from(slot), (U256::from(slot), U256::ZERO))),
+                ),
+            )
+        };
+        // record previous account info
+        bundle_state_builder = bundle_state_builder
+            .revert_account_info(
+                number,
+                address,
+                Some(Some(AccountInfo {
+                    nonce: 1,
+                    balance: U256::from(idx),
+                    ..Default::default()
+                })),
+            )
+            .revert_storage(
+                number,
+                address,
+                Vec::from_iter(slot_range.clone().map(|slot| (U256::from(slot), U256::from(slot)))),
+            );
+    }
+    let bundle = BundleStateWithReceipts::new(
+        bundle_state_builder.build(),
+        Receipts::from_vec(vec![vec![Some(Receipt {
+            tx_type: TxType::EIP1559,
+            success: true,
+            cumulative_gas_used: 400,
+            logs: vec![Log {
+                address: Address::new([0x61; 20]),
+                topics: vec![B256::with_last_byte(3), B256::with_last_byte(4)],
+                data: Bytes::default(),
+            }],
+            #[cfg(feature = "optimism")]
+            deposit_nonce: None,
+            #[cfg(feature = "optimism")]
+            deposit_receipt_version: None,
+        })]]),
+        number,
+    );
+
+    let mut extended = prev_state.clone();
+    extended.extend(bundle.clone());
+    let state_root = bundle_state_root(&extended);
+
+    let mut block = SealedBlock::decode(&mut BLOCK_RLP.as_slice()).unwrap();
+    block.withdrawals = Some(Withdrawals::new(vec![Withdrawal::default()]));
+    let mut header = block.header.clone().unseal();
+    header.number = number;
+    header.state_root = state_root;
+    // parent_hash points to block1 hash
+    header.parent_hash = parent_hash;
+    block.header = header.seal_slow();
+
+    (SealedBlockWithSenders { block, senders: vec![Address::new([0x31; 20])] }, bundle)
+}
+
+/// Block five that points to block 4
+fn block5(
+    number: BlockNumber,
+    parent_hash: B256,
+    prev_state: &BundleStateWithReceipts,
+) -> (SealedBlockWithSenders, BundleStateWithReceipts) {
+    let address_range = 1..=20;
+    let slot_range = 1..=100;
+
+    let mut bundle_state_builder = BundleState::builder(number..=number);
+    for idx in address_range {
+        let address = Address::with_last_byte(idx);
+        // update every even account and recreate every odd only with half of slots
+        bundle_state_builder = bundle_state_builder
+            .state_present_account_info(
+                address,
+                AccountInfo { nonce: 1, balance: U256::from(idx * 2), ..Default::default() },
+            )
+            .state_storage(
+                address,
+                HashMap::from_iter(
+                    slot_range
+                        .clone()
+                        .take(50)
+                        .map(|slot| (U256::from(slot), (U256::from(slot), U256::from(slot * 4)))),
+                ),
+            );
+        bundle_state_builder = if idx % 2 == 0 {
+            bundle_state_builder
+                .revert_account_info(
+                    number,
+                    address,
+                    Some(Some(AccountInfo {
+                        nonce: 1,
+                        balance: U256::from(idx * 2),
+                        ..Default::default()
+                    })),
+                )
+                .revert_storage(
+                    number,
+                    address,
+                    Vec::from_iter(
+                        slot_range.clone().map(|slot| (U256::from(slot), U256::from(slot * 2))),
+                    ),
+                )
+        } else {
+            bundle_state_builder.revert_address(number, address)
+        };
+    }
+    let bundle = BundleStateWithReceipts::new(
+        bundle_state_builder.build(),
+        Receipts::from_vec(vec![vec![Some(Receipt {
+            tx_type: TxType::EIP1559,
+            success: true,
+            cumulative_gas_used: 400,
+            logs: vec![Log {
+                address: Address::new([0x61; 20]),
+                topics: vec![B256::with_last_byte(3), B256::with_last_byte(4)],
+                data: Bytes::default(),
+            }],
+            #[cfg(feature = "optimism")]
+            deposit_nonce: None,
+            #[cfg(feature = "optimism")]
+            deposit_receipt_version: None,
+        })]]),
+        number,
+    );
+
+    let mut extended = prev_state.clone();
+    extended.extend(bundle.clone());
+    let state_root = bundle_state_root(&extended);
+
+    let mut block = SealedBlock::decode(&mut BLOCK_RLP.as_slice()).unwrap();
+    block.withdrawals = Some(Withdrawals::new(vec![Withdrawal::default()]));
+    let mut header = block.header.clone().unseal();
+    header.number = number;
+    header.state_root = state_root;
+    // parent_hash points to block1 hash
+    header.parent_hash = parent_hash;
+    block.header = header.seal_slow();
+
+    (SealedBlockWithSenders { block, senders: vec![Address::new([0x31; 20])] }, bundle)
+}