From 223dde200fb8432ae3eae92a7e4606fff0d533dd Mon Sep 17 00:00:00 2001
From: Roman Krasiuk <rokrassyuk@gmail.com>
Date: Sun, 21 Apr 2024 11:20:55 +0200
Subject: [PATCH] fix(tree): disable cached trie updates for chains with >1
 block (#7753)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>
---
 crates/blockchain-tree/src/blockchain_tree.rs |  4 +--
 crates/blockchain-tree/src/chain.rs           |  4 +--
 crates/storage/provider/src/chain.rs          | 26 ++++---------------
 3 files changed, 9 insertions(+), 25 deletions(-)

diff --git a/crates/blockchain-tree/src/blockchain_tree.rs b/crates/blockchain-tree/src/blockchain_tree.rs
index eb699ff1af63..02bae76bb7b7 100644
--- a/crates/blockchain-tree/src/blockchain_tree.rs
+++ b/crates/blockchain-tree/src/blockchain_tree.rs
@@ -1715,7 +1715,7 @@ mod tests {
         );
         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!(block2_chain.trie_updates().is_none());
 
         assert_eq!(
             tree.make_canonical(block2.hash()).unwrap(),
@@ -1750,7 +1750,7 @@ mod tests {
 
         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!(block5_chain.trie_updates().is_none());
 
         assert_eq!(
             tree.make_canonical(block5.hash()).unwrap(),
diff --git a/crates/blockchain-tree/src/chain.rs b/crates/blockchain-tree/src/chain.rs
index abefae581790..2444cf24a901 100644
--- a/crates/blockchain-tree/src/chain.rs
+++ b/crates/blockchain-tree/src/chain.rs
@@ -282,7 +282,7 @@ impl AppendableChain {
             canonical_fork,
         };
 
-        let (block_state, trie_updates) = Self::validate_and_execute(
+        let (block_state, _) = Self::validate_and_execute(
             block.clone(),
             parent_block,
             bundle_state_data,
@@ -291,7 +291,7 @@ impl AppendableChain {
             block_validation_kind,
         )?;
         // extend the state.
-        self.chain.append_block(block, block_state, trie_updates);
+        self.chain.append_block(block, block_state);
 
         Ok(())
     }
diff --git a/crates/storage/provider/src/chain.rs b/crates/storage/provider/src/chain.rs
index 4114ba097006..5acd84599748 100644
--- a/crates/storage/provider/src/chain.rs
+++ b/crates/storage/provider/src/chain.rs
@@ -26,7 +26,8 @@ pub struct Chain {
     /// This state also contains the individual changes that lead to the current state.
     state: BundleStateWithReceipts,
     /// State trie updates after block is added to the chain.
-    /// NOTE: Currently, trie updates are present only if the block extends canonical chain.
+    /// NOTE: Currently, trie updates are present only for
+    /// single-block chains that extend the canonical chain.
     trie_updates: Option<TrieUpdates>,
 }
 
@@ -210,15 +211,10 @@ impl Chain {
 
     /// Append a single block with state to the chain.
     /// This method assumes that blocks attachment to the chain has already been validated.
-    pub fn append_block(
-        &mut self,
-        block: SealedBlockWithSenders,
-        state: BundleStateWithReceipts,
-        trie_updates: Option<TrieUpdates>,
-    ) {
+    pub fn append_block(&mut self, block: SealedBlockWithSenders, state: BundleStateWithReceipts) {
         self.blocks.insert(block.number, block);
         self.state.extend(state);
-        self.append_trie_updates(trie_updates);
+        self.trie_updates.take(); // reset
     }
 
     /// Merge two chains by appending the given chain into the current one.
@@ -238,23 +234,11 @@ impl Chain {
         // Insert blocks from other chain
         self.blocks.extend(other.blocks);
         self.state.extend(other.state);
-        self.append_trie_updates(other.trie_updates);
+        self.trie_updates.take(); // reset
 
         Ok(())
     }
 
-    /// Append trie updates.
-    /// If existing or incoming trie updates are not set, reset as neither is valid anymore.
-    fn append_trie_updates(&mut self, other_trie_updates: Option<TrieUpdates>) {
-        if let Some((trie_updates, other)) = self.trie_updates.as_mut().zip(other_trie_updates) {
-            // Extend trie updates.
-            trie_updates.extend(other);
-        } else {
-            // Reset trie updates as they are no longer valid.
-            self.trie_updates.take();
-        }
-    }
-
     /// Split this chain at the given block.
     ///
     /// The given block will be the last block in the first returned chain.