From 5eef3c5fc6c3ad2ecfbc2ebdcd45a0838bc77d4c Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Tue, 18 Apr 2023 16:54:53 +0200
Subject: [PATCH] chore: add ChainBlocks type (#2297)

Co-authored-by: rakita <rakita@users.noreply.github.com>
---
 .../src/blockchain_tree/mod.rs                |  2 +-
 crates/storage/provider/src/chain.rs          | 53 ++++++++++++++++---
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/crates/blockchain-tree/src/blockchain_tree/mod.rs b/crates/blockchain-tree/src/blockchain_tree/mod.rs
index 0b2bbfceb0..637581b6ad 100644
--- a/crates/blockchain-tree/src/blockchain_tree/mod.rs
+++ b/crates/blockchain-tree/src/blockchain_tree/mod.rs
@@ -625,7 +625,7 @@ impl<DB: Database, C: Consensus, EF: ExecutorFactory> BlockchainTree<DB, C, EF>
 
         let (blocks, state) = chain.into_inner();
 
-        tx.append_blocks_with_post_state(blocks.into_values().collect(), state)
+        tx.append_blocks_with_post_state(blocks.into_blocks().collect(), state)
             .map_err(|e| ExecError::CanonicalCommit { inner: e.to_string() })?;
 
         tx.commit()?;
diff --git a/crates/storage/provider/src/chain.rs b/crates/storage/provider/src/chain.rs
index c57a510e18..f754c46c5f 100644
--- a/crates/storage/provider/src/chain.rs
+++ b/crates/storage/provider/src/chain.rs
@@ -16,11 +16,11 @@ use std::collections::BTreeMap;
 /// Used inside the BlockchainTree.
 #[derive(Clone, Debug, Default, PartialEq, Eq)]
 pub struct Chain {
-    /// The state of accounts after execution of the blocks in this chain.
+    /// The state of accounts after execution of the blocks in this chain (tip of the chain).
     ///
     /// This state also contains the individual changes that lead to the current state.
     pub state: PostState,
-    /// The blocks in this chain.
+    /// All blocks in this chain.
     pub blocks: BTreeMap<BlockNumber, SealedBlockWithSenders>,
     /// A mapping of each block number in the chain to the highest transition ID in the chain's
     /// state after execution of the block.
@@ -60,12 +60,12 @@ impl Chain {
 
     /// Return post state of the block at the `block_number` or None if block is not known
     pub fn state_at_block(&self, block_number: BlockNumber) -> Option<PostState> {
-        let mut state = self.state.clone();
         if self.tip().number == block_number {
-            return Some(state)
+            return Some(self.state.clone())
         }
 
         if let Some(&transition_id) = self.block_transitions.get(&block_number) {
+            let mut state = self.state.clone();
             state.revert_to(transition_id);
             return Some(state)
         }
@@ -73,9 +73,10 @@ impl Chain {
         None
     }
 
-    /// Destructure the chain into its inner components, the blocks and the state.
-    pub fn into_inner(self) -> (BTreeMap<BlockNumber, SealedBlockWithSenders>, PostState) {
-        (self.blocks, self.state)
+    /// Destructure the chain into its inner components, the blocks and the state at the tip of the
+    /// chain.
+    pub fn into_inner(self) -> (ChainBlocks, PostState) {
+        (ChainBlocks { blocks: self.blocks }, self.state)
     }
 
     /// Get the block at which this chain forked.
@@ -228,6 +229,44 @@ impl Chain {
     }
 }
 
+/// All blocks in the chain
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub struct ChainBlocks {
+    blocks: BTreeMap<BlockNumber, SealedBlockWithSenders>,
+}
+
+impl ChainBlocks {
+    /// Creates a consuming iterator over all blocks in the chain with increasing block number.
+    ///
+    /// Note: this always yields at least one block.
+    pub fn into_blocks(self) -> impl Iterator<Item = SealedBlockWithSenders> {
+        self.blocks.into_values()
+    }
+
+    /// Creates an iterator over all blocks in the chain with increasing block number.
+    pub fn iter(&self) -> impl Iterator<Item = (&BlockNumber, &SealedBlockWithSenders)> {
+        self.blocks.iter()
+    }
+
+    /// Get the tip of the chain.
+    ///
+    /// # Note
+    ///
+    /// Chains always have at least one block.
+    pub fn tip(&self) -> &SealedBlockWithSenders {
+        self.blocks.last_key_value().expect("Chain should have at least one block").1
+    }
+}
+
+impl IntoIterator for ChainBlocks {
+    type Item = (BlockNumber, SealedBlockWithSenders);
+    type IntoIter = std::collections::btree_map::IntoIter<BlockNumber, SealedBlockWithSenders>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.blocks.into_iter()
+    }
+}
+
 /// Used to hold receipts and their attachment.
 #[derive(Default, Clone, Debug)]
 pub struct BlockReceipts {