From ed9f6558777ae4248df9da073c76a98a33387643 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Thu, 6 Jun 2019 00:27:17 -0400
Subject: [PATCH] Update `BeaconChain` iters

---
 beacon_node/beacon_chain/src/beacon_chain.rs | 17 ++++++++---
 beacon_node/beacon_chain/src/iter.rs         | 31 +++++++++++++++++---
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs
index 13aad8112..d20813969 100644
--- a/beacon_node/beacon_chain/src/beacon_chain.rs
+++ b/beacon_node/beacon_chain/src/beacon_chain.rs
@@ -1,8 +1,8 @@
 use crate::checkpoint::CheckPoint;
 use crate::errors::{BeaconChainError as Error, BlockProductionError};
+use crate::iter::{BlockIterator, BlockRootsIterator};
 use crate::metrics::Metrics;
 use crate::persisted_beacon_chain::{PersistedBeaconChain, BEACON_CHAIN_DB_KEY};
-use crate::iter::BlockRootsIterator;
 use fork_choice::{ForkChoice, ForkChoiceError};
 use log::{debug, trace};
 use operation_pool::DepositInsertStatus;
@@ -226,15 +226,24 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
 
         Ok(headers?)
     }
+    /// Iterate in reverse (highest to lowest slot) through all blocks from the block at `slot`
+    /// through to the genesis block.
+    ///
+    /// Returns `None` for headers prior to genesis or when there is an error reading from `Store`.
+    ///
+    /// Contains duplicate headers when skip slots are encountered.
+    pub fn rev_iter_blocks(&self, slot: Slot) -> BlockIterator<T::EthSpec, T::Store> {
+        BlockIterator::new(self.store.clone(), self.state.read().clone(), slot)
+    }
 
-    /// Iterate in reverse through all block roots starting from the current state, through to
+    /// Iterates in reverse (highest to lowest slot) through all block roots from `slot` through to
     /// genesis.
     ///
     /// Returns `None` for roots prior to genesis or when there is an error reading from `Store`.
     ///
     /// Contains duplicate roots when skip slots are encountered.
-    pub fn iter_block_roots(&self) -> BlockRootsIterator<T::EthSpec, T::Store> {
-        BlockRootsIterator::from_state(self.store.clone(), self.state.read().clone())
+    pub fn rev_iter_block_roots(&self, slot: Slot) -> BlockRootsIterator<T::EthSpec, T::Store> {
+        BlockRootsIterator::new(self.store.clone(), self.state.read().clone(), slot)
     }
 
     /*
diff --git a/beacon_node/beacon_chain/src/iter.rs b/beacon_node/beacon_chain/src/iter.rs
index da81a6582..da939be2a 100644
--- a/beacon_node/beacon_chain/src/iter.rs
+++ b/beacon_node/beacon_chain/src/iter.rs
@@ -1,6 +1,29 @@
 use std::sync::Arc;
 use store::Store;
-use types::{BeaconState, BeaconStateError, EthSpec, Hash256, Slot};
+use types::{BeaconBlock, BeaconState, BeaconStateError, EthSpec, Hash256, Slot};
+
+/// Extends `BlockRootsIterator`, returning `BeaconBlock` instances, instead of their roots.
+pub struct BlockIterator<T: EthSpec, U> {
+    roots: BlockRootsIterator<T, U>,
+}
+
+impl<T: EthSpec, U: Store> BlockIterator<T, U> {
+    /// Create a new iterator over all blocks in the given `beacon_state` and prior states.
+    pub fn new(store: Arc<U>, beacon_state: BeaconState<T>, start_slot: Slot) -> Self {
+        Self {
+            roots: BlockRootsIterator::new(store, beacon_state, start_slot),
+        }
+    }
+}
+
+impl<T: EthSpec, U: Store> Iterator for BlockIterator<T, U> {
+    type Item = BeaconBlock;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let root = self.roots.next()?;
+        self.roots.store.get(&root).ok()?
+    }
+}
 
 /// Iterates backwards through block roots.
 ///
@@ -17,9 +40,9 @@ pub struct BlockRootsIterator<T: EthSpec, U> {
 
 impl<T: EthSpec, U: Store> BlockRootsIterator<T, U> {
     /// Create a new iterator over all block roots in the given `beacon_state` and prior states.
-    pub fn from_state(store: Arc<U>, beacon_state: BeaconState<T>) -> Self {
+    pub fn new(store: Arc<U>, beacon_state: BeaconState<T>, start_slot: Slot) -> Self {
         Self {
-            slot: beacon_state.slot,
+            slot: start_slot,
             beacon_state,
             store,
         }
@@ -95,7 +118,7 @@ mod test {
         state_b.latest_state_roots[0] = state_a_root;
         store.put(&state_a_root, &state_a).unwrap();
 
-        let iter = BlockRootsIterator::from_state(store.clone(), state_b.clone());
+        let iter = BlockRootsIterator::new(store.clone(), state_b.clone(), state_b.slot);
         let mut collected: Vec<Hash256> = iter.collect();
         collected.reverse();