Implement tree states & hierarchical state DB

beacon_node/store/src/forwards_iter.rs

@@ -1,29 +1,34 @@
-use crate::chunked_iter::ChunkedVectorIter;
-use crate::chunked_vector::{BlockRoots, Field, StateRoots};
 use crate::errors::{Error, Result};
 use crate::iter::{BlockRootsIterator, StateRootsIterator};
-use crate::{HotColdDB, ItemStore};
+use crate::{ColumnIter, DBColumn, HotColdDB, ItemStore};
 use itertools::process_results;
-use types::{BeaconState, ChainSpec, EthSpec, Hash256, Slot};
+use std::marker::PhantomData;
+use types::{BeaconState, EthSpec, Hash256, Slot};
 
 pub type HybridForwardsBlockRootsIterator<'a, E, Hot, Cold> =
-    HybridForwardsIterator<'a, E, BlockRoots, Hot, Cold>;
+    HybridForwardsIterator<'a, E, Hot, Cold>;
 pub type HybridForwardsStateRootsIterator<'a, E, Hot, Cold> =
-    HybridForwardsIterator<'a, E, StateRoots, Hot, Cold>;
+    HybridForwardsIterator<'a, E, Hot, Cold>;
 
-/// Trait unifying `BlockRoots` and `StateRoots` for forward iteration.
-pub trait Root<E: EthSpec>: Field<E, Value = Hash256> {
-    fn simple_forwards_iterator<Hot: ItemStore<E>, Cold: ItemStore<E>>(
-        store: &HotColdDB<E, Hot, Cold>,
+impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold> {
+    pub fn simple_forwards_iterator(
+        &self,
+        column: DBColumn,
         start_slot: Slot,
         end_state: BeaconState<E>,
         end_root: Hash256,
-    ) -> Result<SimpleForwardsIterator>;
-}
+    ) -> Result<SimpleForwardsIterator> {
+        if column == DBColumn::BeaconBlockRoots {
+            self.forwards_iter_block_roots_using_state(start_slot, end_state, end_root)
+        } else if column == DBColumn::BeaconStateRoots {
+            self.forwards_iter_state_roots_using_state(start_slot, end_state, end_root)
+        } else {
+            panic!("FIXME(sproul): better error")
+        }
+    }
 
-impl<E: EthSpec> Root<E> for BlockRoots {
-    fn simple_forwards_iterator<Hot: ItemStore<E>, Cold: ItemStore<E>>(
-        store: &HotColdDB<E, Hot, Cold>,
+    pub fn forwards_iter_block_roots_using_state(
+        &self,
         start_slot: Slot,
         end_state: BeaconState<E>,
         end_block_root: Hash256,
@@ -31,7 +36,7 @@ impl<E: EthSpec> Root<E> for BlockRoots {
         // Iterate backwards from the end state, stopping at the start slot.
         let values = process_results(
             std::iter::once(Ok((end_block_root, end_state.slot())))
-                .chain(BlockRootsIterator::owned(store, end_state)),
+                .chain(BlockRootsIterator::owned(self, end_state)),
             |iter| {
                 iter.take_while(|(_, slot)| *slot >= start_slot)
                     .collect::<Vec<_>>()
@@ -39,11 +44,9 @@ impl<E: EthSpec> Root<E> for BlockRoots {
         )?;
         Ok(SimpleForwardsIterator { values })
     }
-}
 
-impl<E: EthSpec> Root<E> for StateRoots {
-    fn simple_forwards_iterator<Hot: ItemStore<E>, Cold: ItemStore<E>>(
-        store: &HotColdDB<E, Hot, Cold>,
+    pub fn forwards_iter_state_roots_using_state(
+        &self,
         start_slot: Slot,
         end_state: BeaconState<E>,
         end_state_root: Hash256,
@@ -51,7 +54,7 @@ impl<E: EthSpec> Root<E> for StateRoots {
         // 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)),
+                .chain(StateRootsIterator::owned(self, end_state)),
             |iter| {
                 iter.take_while(|(_, slot)| *slot >= start_slot)
                     .collect::<Vec<_>>()
@@ -62,40 +65,62 @@ impl<E: EthSpec> Root<E> for StateRoots {
 }
 
 /// Forwards root iterator that makes use of a flat field table in the freezer DB.
-pub struct FrozenForwardsIterator<'a, E: EthSpec, F: Root<E>, Hot: ItemStore<E>, Cold: ItemStore<E>>
-{
-    inner: ChunkedVectorIter<'a, F, E, Hot, Cold>,
+pub struct FrozenForwardsIterator<'a, E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> {
+    inner: ColumnIter<'a, Vec<u8>>,
+    limit: Slot,
+    finished: bool,
+    _phantom: PhantomData<(E, Hot, Cold)>,
 }
 
-impl<'a, E: EthSpec, F: Root<E>, Hot: ItemStore<E>, Cold: ItemStore<E>>
-    FrozenForwardsIterator<'a, E, F, Hot, Cold>
+impl<'a, E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>>
+    FrozenForwardsIterator<'a, E, Hot, Cold>
 {
+    /// `end_slot` is EXCLUSIVE here.
     pub fn new(
         store: &'a HotColdDB<E, Hot, Cold>,
+        column: DBColumn,
         start_slot: Slot,
-        last_restore_point_slot: Slot,
-        spec: &ChainSpec,
+        end_slot: Slot,
     ) -> Self {
+        if column != DBColumn::BeaconBlockRoots && column != DBColumn::BeaconStateRoots {
+            panic!("FIXME(sproul): bad column error");
+        }
+        let start = start_slot.as_u64().to_be_bytes();
         Self {
-            inner: ChunkedVectorIter::new(
-                store,
-                start_slot.as_usize(),
-                last_restore_point_slot,
-                spec,
-            ),
+            inner: store.cold_db.iter_column_from(column, &start),
+            limit: end_slot,
+            finished: false,
+            _phantom: PhantomData,
         }
     }
 }
 
-impl<'a, E: EthSpec, F: Root<E>, Hot: ItemStore<E>, Cold: ItemStore<E>> Iterator
-    for FrozenForwardsIterator<'a, E, F, Hot, Cold>
+impl<'a, E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> Iterator
+    for FrozenForwardsIterator<'a, E, Hot, Cold>
 {
-    type Item = (Hash256, Slot);
+    type Item = Result<(Hash256, Slot)>;
 
     fn next(&mut self) -> Option<Self::Item> {
+        if self.finished {
+            return None;
+        }
+
         self.inner
-            .next()
-            .map(|(slot, root)| (root, Slot::from(slot)))
+            .next()?
+            .and_then(|(slot_bytes, root_bytes)| {
+                if slot_bytes.len() != 8 || root_bytes.len() != 32 {
+                    panic!("FIXME(sproul): put an error here")
+                } else {
+                    let slot = Slot::new(u64::from_be_bytes(slot_bytes.try_into().unwrap()));
+                    let root = Hash256::from_slice(&root_bytes);
+
+                    if slot + 1 == self.limit {
+                        self.finished = true;
+                    }
+                    Ok(Some((root, slot)))
+                }
+            })
+            .transpose()
     }
 }
 
@@ -115,24 +140,27 @@ impl Iterator for SimpleForwardsIterator {
 }
 
 /// Fusion of the above two approaches to forwards iteration. Fast and efficient.
-pub enum HybridForwardsIterator<'a, E: EthSpec, F: Root<E>, Hot: ItemStore<E>, Cold: ItemStore<E>> {
+pub enum HybridForwardsIterator<'a, E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> {
     PreFinalization {
-        iter: Box<FrozenForwardsIterator<'a, E, F, Hot, Cold>>,
+        iter: Box<FrozenForwardsIterator<'a, E, Hot, Cold>>,
+        store: &'a HotColdDB<E, Hot, Cold>,
         /// Data required by the `PostFinalization` iterator when we get to it.
         continuation_data: Option<Box<(BeaconState<E>, Hash256)>>,
+        column: DBColumn,
     },
     PostFinalizationLazy {
         continuation_data: Option<Box<(BeaconState<E>, Hash256)>>,
         store: &'a HotColdDB<E, Hot, Cold>,
         start_slot: Slot,
+        column: DBColumn,
     },
     PostFinalization {
         iter: SimpleForwardsIterator,
     },
 }
 
-impl<'a, E: EthSpec, F: Root<E>, Hot: ItemStore<E>, Cold: ItemStore<E>>
-    HybridForwardsIterator<'a, E, F, Hot, Cold>
+impl<'a, E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>>
+    HybridForwardsIterator<'a, E, Hot, Cold>
 {
     /// Construct a new hybrid iterator.
     ///
@@ -148,41 +176,41 @@ impl<'a, E: EthSpec, F: Root<E>, Hot: ItemStore<E>, Cold: ItemStore<E>>
     /// function may block for some time while `get_state` runs.
     pub fn new(
         store: &'a HotColdDB<E, Hot, Cold>,
+        column: DBColumn,
         start_slot: Slot,
         end_slot: Option<Slot>,
         get_state: impl FnOnce() -> (BeaconState<E>, Hash256),
-        spec: &ChainSpec,
     ) -> Result<Self> {
         use HybridForwardsIterator::*;
 
-        let latest_restore_point_slot = store.get_latest_restore_point_slot();
+        // FIXME(sproul): consider whether this is 100% correct
+        let split_slot = store.get_split_slot();
 
-        let result = if start_slot < latest_restore_point_slot {
+        let result = if start_slot < split_slot {
             let iter = Box::new(FrozenForwardsIterator::new(
-                store,
-                start_slot,
-                latest_restore_point_slot,
-                spec,
+                store, column, start_slot, split_slot,
             ));
 
             // No continuation data is needed if the forwards iterator plans to halt before
             // `end_slot`. If it tries to continue further a `NoContinuationData` error will be
             // returned.
-            let continuation_data =
-                if end_slot.map_or(false, |end_slot| end_slot < latest_restore_point_slot) {
-                    None
-                } else {
-                    Some(Box::new(get_state()))
-                };
+            let continuation_data = if end_slot.map_or(false, |end_slot| end_slot < split_slot) {
+                None
+            } else {
+                Some(Box::new(get_state()))
+            };
             PreFinalization {
                 iter,
+                store,
                 continuation_data,
+                column,
             }
         } else {
             PostFinalizationLazy {
                 continuation_data: Some(Box::new(get_state())),
                 store,
                 start_slot,
+                column,
             }
         };
 
@@ -195,22 +223,24 @@ impl<'a, E: EthSpec, F: Root<E>, Hot: ItemStore<E>, Cold: ItemStore<E>>
         match self {
             PreFinalization {
                 iter,
+                store,
                 continuation_data,
+                column,
             } => {
                 match iter.next() {
-                    Some(x) => Ok(Some(x)),
+                    Some(x) => x.map(Some),
                     // Once the pre-finalization iterator is consumed, transition
                     // to a post-finalization iterator beginning from the last slot
                     // of the pre iterator.
                     None => {
                         let continuation_data = continuation_data.take();
-                        let store = iter.inner.store;
-                        let start_slot = Slot::from(iter.inner.end_vindex);
+                        let start_slot = Slot::from(iter.limit);
 
                         *self = PostFinalizationLazy {
                             continuation_data,
                             store,
                             start_slot,
+                            column: *column,
                         };
 
                         self.do_next()
@@ -221,11 +251,17 @@ impl<'a, E: EthSpec, F: Root<E>, Hot: ItemStore<E>, Cold: ItemStore<E>>
                 continuation_data,
                 store,
                 start_slot,
+                column,
             } => {
                 let (end_state, end_root) =
                     *continuation_data.take().ok_or(Error::NoContinuationData)?;
                 *self = PostFinalization {
-                    iter: F::simple_forwards_iterator(store, *start_slot, end_state, end_root)?,
+                    iter: store.simple_forwards_iterator(
+                        *column,
+                        *start_slot,
+                        end_state,
+                        end_root,
+                    )?,
                 };
                 self.do_next()
             }
@@ -234,8 +270,8 @@ impl<'a, E: EthSpec, F: Root<E>, Hot: ItemStore<E>, Cold: ItemStore<E>>
     }
 }
 
-impl<'a, E: EthSpec, F: Root<E>, Hot: ItemStore<E>, Cold: ItemStore<E>> Iterator
-    for HybridForwardsIterator<'a, E, F, Hot, Cold>
+impl<'a, E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> Iterator
+    for HybridForwardsIterator<'a, E, Hot, Cold>
 {
     type Item = Result<(Hash256, Slot)>;