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 BeaconChain { 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 { + 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 { - BlockRootsIterator::from_state(self.store.clone(), self.state.read().clone()) + pub fn rev_iter_block_roots(&self, slot: Slot) -> BlockRootsIterator { + 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 { + roots: BlockRootsIterator, +} + +impl BlockIterator { + /// Create a new iterator over all blocks in the given `beacon_state` and prior states. + pub fn new(store: Arc, beacon_state: BeaconState, start_slot: Slot) -> Self { + Self { + roots: BlockRootsIterator::new(store, beacon_state, start_slot), + } + } +} + +impl Iterator for BlockIterator { + type Item = BeaconBlock; + + fn next(&mut self) -> Option { + let root = self.roots.next()?; + self.roots.store.get(&root).ok()? + } +} /// Iterates backwards through block roots. /// @@ -17,9 +40,9 @@ pub struct BlockRootsIterator { impl BlockRootsIterator { /// Create a new iterator over all block roots in the given `beacon_state` and prior states. - pub fn from_state(store: Arc, beacon_state: BeaconState) -> Self { + pub fn new(store: Arc, beacon_state: BeaconState, 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 = iter.collect(); collected.reverse();