From 6757ab81902b4c0e8a5ec51eb7d9ed10ec328367 Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Fri, 27 Sep 2024 14:47:00 +0200
Subject: [PATCH] fix: make canonical_chain atomic and canonical (#11283)

---
 crates/chain-state/src/in_memory.rs | 53 ++++++++++++-----------------
 1 file changed, 21 insertions(+), 32 deletions(-)

diff --git a/crates/chain-state/src/in_memory.rs b/crates/chain-state/src/in_memory.rs
index 029666f6ac..fc142dd03a 100644
--- a/crates/chain-state/src/in_memory.rs
+++ b/crates/chain-state/src/in_memory.rs
@@ -512,23 +512,12 @@ impl CanonicalInMemoryState {
         MemoryOverlayStateProvider::new(historical, in_memory)
     }
 
-    /// Returns an iterator over all canonical blocks in the in-memory state, from newest to oldest
-    /// (highest to lowest).
+    /// Returns an iterator over all __canonical blocks__ in the in-memory state, from newest to
+    /// oldest (highest to lowest).
+    ///
+    /// This iterator contains a snapshot of the in-memory state at the time of the call.
     pub fn canonical_chain(&self) -> impl Iterator<Item = Arc<BlockState>> {
-        let pending = self.inner.in_memory_state.pending.borrow().clone();
-        let head = self.inner.in_memory_state.head_state();
-
-        // this clone is cheap because we only expect to keep in memory a few
-        // blocks and all of them are Arcs.
-        let blocks = self.inner.in_memory_state.blocks.read().clone();
-
-        std::iter::once(pending).filter_map(|p| p.map(Arc::new)).chain(std::iter::successors(
-            head,
-            move |state| {
-                let parent_hash = state.block().block().parent_hash;
-                blocks.get(&parent_hash).cloned()
-            },
-        ))
+        self.inner.in_memory_state.head_state().into_iter().flat_map(|head| head.iter())
     }
 
     /// Returns a `TransactionSigned` for the given `TxHash` if found.
@@ -693,6 +682,13 @@ impl BlockState {
     pub fn append_parent_chain<'a>(&'a self, chain: &mut Vec<&'a Self>) {
         chain.extend(self.parent_state_chain());
     }
+
+    /// Returns an iterator over the atomically captured chain of in memory blocks.
+    ///
+    /// This yields the blocks from newest to oldest (highest to lowest).
+    pub fn iter(self: Arc<Self>) -> impl Iterator<Item = Arc<Self>> {
+        std::iter::successors(Some(self), |state| state.parent.clone())
+    }
 }
 
 /// Represents an executed block stored in-memory.
@@ -1267,20 +1263,17 @@ mod tests {
 
     #[test]
     fn test_canonical_in_memory_state_canonical_chain_multiple_blocks() {
-        let mut blocks = HashMap::default();
-        let mut numbers = BTreeMap::new();
         let mut parent_hash = B256::random();
         let mut block_builder = TestBlockBuilder::default();
+        let state = CanonicalInMemoryState::empty();
 
         for i in 1..=3 {
             let block = block_builder.get_executed_block_with_number(i, parent_hash);
             let hash = block.block().hash();
-            blocks.insert(hash, Arc::new(BlockState::new(block.clone())));
-            numbers.insert(i, hash);
+            state.update_blocks(Some(block), None);
             parent_hash = hash;
         }
 
-        let state = CanonicalInMemoryState::new(blocks, numbers, None, None);
         let chain: Vec<_> = state.canonical_chain().collect();
 
         assert_eq!(chain.len(), 3);
@@ -1289,31 +1282,27 @@ mod tests {
         assert_eq!(chain[2].number(), 1);
     }
 
+    // ensures the pending block is not part of the canonical chain
     #[test]
     fn test_canonical_in_memory_state_canonical_chain_with_pending_block() {
-        let mut blocks = HashMap::default();
-        let mut numbers = BTreeMap::new();
         let mut parent_hash = B256::random();
         let mut block_builder = TestBlockBuilder::default();
+        let state = CanonicalInMemoryState::empty();
 
         for i in 1..=2 {
             let block = block_builder.get_executed_block_with_number(i, parent_hash);
             let hash = block.block().hash();
-            blocks.insert(hash, Arc::new(BlockState::new(block.clone())));
-            numbers.insert(i, hash);
+            state.update_blocks(Some(block), None);
             parent_hash = hash;
         }
 
         let pending_block = block_builder.get_executed_block_with_number(3, parent_hash);
-        let pending_state = BlockState::new(pending_block);
-
-        let state = CanonicalInMemoryState::new(blocks, numbers, Some(pending_state), None);
+        state.set_pending_block(pending_block);
         let chain: Vec<_> = state.canonical_chain().collect();
 
-        assert_eq!(chain.len(), 3);
-        assert_eq!(chain[0].number(), 3);
-        assert_eq!(chain[1].number(), 2);
-        assert_eq!(chain[2].number(), 1);
+        assert_eq!(chain.len(), 2);
+        assert_eq!(chain[0].number(), 2);
+        assert_eq!(chain[1].number(), 1);
     }
 
     #[test]