Use forwards iterator for state root lookups (#2422)

## Issue Addressed

#2377 

## Proposed Changes

Implement the same code used for block root lookups (from #2376) to state root lookups in order to improve performance and reduce associated memory spikes (e.g. from certain HTTP API requests).

## Additional Changes

- Tests using `rev_iter_state_roots` and `rev_iter_block_roots` have been refactored to use their `forwards` versions instead.
- The `rev_iter_state_roots` and `rev_iter_block_roots` functions are now unused and have been removed.
- The `state_at_slot` function has been changed to use the `forwards` iterator.

## Additional Info

- Some tests still need to be refactored to use their `forwards_iter` versions. These tests start their iteration from a specific beacon state and thus use the `rev_iter_state_roots_from` and `rev_iter_block_roots_from` functions. If they can be refactored, those functions can also be removed.

beacon_node/store/src/forwards_iter.rs

@@ -1,7 +1,7 @@
 use crate::chunked_iter::ChunkedVectorIter;
-use crate::chunked_vector::BlockRoots;
+use crate::chunked_vector::{BlockRoots, StateRoots};
 use crate::errors::{Error, Result};
-use crate::iter::BlockRootsIterator;
+use crate::iter::{BlockRootsIterator, StateRootsIterator};
 use crate::{HotColdDB, ItemStore};
 use itertools::process_results;
 use std::sync::Arc;
@@ -172,3 +172,169 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> Iterator
         self.do_next().transpose()
     }
 }
+
+/// Forwards state roots iterator that makes use of the `state_roots` table in the freezer DB.
+pub struct FrozenForwardsStateRootsIterator<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> {
+    inner: ChunkedVectorIter<StateRoots, E, Hot, Cold>,
+}
+
+/// Forwards state roots iterator that reverses a backwards iterator (only good for short ranges).
+pub struct SimpleForwardsStateRootsIterator {
+    // Values from the backwards iterator (in slot descending order)
+    values: Vec<(Hash256, Slot)>,
+}
+
+/// Fusion of the above two approaches to forwards iteration. Fast and efficient.
+pub enum HybridForwardsStateRootsIterator<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> {
+    PreFinalization {
+        iter: Box<FrozenForwardsStateRootsIterator<E, Hot, Cold>>,
+        /// Data required by the `PostFinalization` iterator when we get to it.
+        continuation_data: Box<Option<(BeaconState<E>, Hash256)>>,
+    },
+    PostFinalization {
+        iter: SimpleForwardsStateRootsIterator,
+    },
+}
+
+impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>>
+    FrozenForwardsStateRootsIterator<E, Hot, Cold>
+{
+    pub fn new(
+        store: Arc<HotColdDB<E, Hot, Cold>>,
+        start_slot: Slot,
+        last_restore_point_slot: Slot,
+        spec: &ChainSpec,
+    ) -> Self {
+        Self {
+            inner: ChunkedVectorIter::new(
+                store,
+                start_slot.as_usize(),
+                last_restore_point_slot,
+                spec,
+            ),
+        }
+    }
+}
+
+impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> Iterator
+    for FrozenForwardsStateRootsIterator<E, Hot, Cold>
+{
+    type Item = (Hash256, Slot);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.inner
+            .next()
+            .map(|(slot, state_hash)| (state_hash, Slot::from(slot)))
+    }
+}
+
+impl SimpleForwardsStateRootsIterator {
+    pub fn new<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>>(
+        store: Arc<HotColdDB<E, Hot, Cold>>,
+        start_slot: Slot,
+        end_state: BeaconState<E>,
+        end_state_root: Hash256,
+    ) -> Result<Self> {
+        // Iterate backwards from the end state, stopping at the start slot.
+        let values = process_results(
+            std::iter::once(Ok((end_state_root, end_state.slot)))
+                .chain(StateRootsIterator::owned(store, end_state)),
+            |iter| {
+                iter.take_while(|(_, slot)| *slot >= start_slot)
+                    .collect::<Vec<_>>()
+            },
+        )?;
+        Ok(Self { values })
+    }
+}
+
+impl Iterator for SimpleForwardsStateRootsIterator {
+    type Item = Result<(Hash256, Slot)>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // Pop from the end of the vector to get the state roots in slot-ascending order.
+        Ok(self.values.pop()).transpose()
+    }
+}
+
+impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>>
+    HybridForwardsStateRootsIterator<E, Hot, Cold>
+{
+    pub fn new(
+        store: Arc<HotColdDB<E, Hot, Cold>>,
+        start_slot: Slot,
+        end_state: BeaconState<E>,
+        end_state_root: Hash256,
+        spec: &ChainSpec,
+    ) -> Result<Self> {
+        use HybridForwardsStateRootsIterator::*;
+
+        let latest_restore_point_slot = store.get_latest_restore_point_slot();
+
+        let result = if start_slot < latest_restore_point_slot {
+            PreFinalization {
+                iter: Box::new(FrozenForwardsStateRootsIterator::new(
+                    store,
+                    start_slot,
+                    latest_restore_point_slot,
+                    spec,
+                )),
+                continuation_data: Box::new(Some((end_state, end_state_root))),
+            }
+        } else {
+            PostFinalization {
+                iter: SimpleForwardsStateRootsIterator::new(
+                    store,
+                    start_slot,
+                    end_state,
+                    end_state_root,
+                )?,
+            }
+        };
+
+        Ok(result)
+    }
+
+    fn do_next(&mut self) -> Result<Option<(Hash256, Slot)>> {
+        use HybridForwardsStateRootsIterator::*;
+
+        match self {
+            PreFinalization {
+                iter,
+                continuation_data,
+            } => {
+                match iter.next() {
+                    Some(x) => Ok(Some(x)),
+                    // Once the pre-finalization iterator is consumed, transition
+                    // to a post-finalization iterator beginning from the last slot
+                    // of the pre iterator.
+                    None => {
+                        let (end_state, end_state_root) =
+                            continuation_data.take().ok_or(Error::NoContinuationData)?;
+
+                        *self = PostFinalization {
+                            iter: SimpleForwardsStateRootsIterator::new(
+                                iter.inner.store.clone(),
+                                Slot::from(iter.inner.end_vindex),
+                                end_state,
+                                end_state_root,
+                            )?,
+                        };
+                        self.do_next()
+                    }
+                }
+            }
+            PostFinalization { iter } => iter.next().transpose(),
+        }
+    }
+}
+
+impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> Iterator
+    for HybridForwardsStateRootsIterator<E, Hot, Cold>
+{
+    type Item = Result<(Hash256, Slot)>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.do_next().transpose()
+    }
+}

beacon_node/store/src/hot_cold_store.rs

@@ -2,7 +2,7 @@ use crate::chunked_vector::{
     store_updated_vector, BlockRoots, HistoricalRoots, RandaoMixes, StateRoots,
 };
 use crate::config::{OnDiskStoreConfig, StoreConfig};
-use crate::forwards_iter::HybridForwardsBlockRootsIterator;
+use crate::forwards_iter::{HybridForwardsBlockRootsIterator, HybridForwardsStateRootsIterator};
 use crate::impls::beacon_state::{get_full_state, store_full_state};
 use crate::iter::{ParentRootBlockIterator, StateRootsIterator};
 use crate::leveldb_store::BytesKey;
@@ -393,6 +393,16 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
         HybridForwardsBlockRootsIterator::new(store, start_slot, end_state, end_block_root, spec)
     }
 
+    pub fn forwards_state_roots_iterator(
+        store: Arc<Self>,
+        start_slot: Slot,
+        end_state_root: Hash256,
+        end_state: BeaconState<E>,
+        spec: &ChainSpec,
+    ) -> Result<impl Iterator<Item = Result<(Hash256, Slot), Error>>, Error> {
+        HybridForwardsStateRootsIterator::new(store, start_slot, end_state, end_state_root, spec)
+    }
+
     /// Load an epoch boundary state by using the hot state summary look-up.
     ///
     /// Will fall back to the cold DB if a hot state summary is not found.