Add state roots iter to store

beacon_node/store/src/iter.rs

@@ -1,22 +1,73 @@
 use crate::Store;
+use std::borrow::Cow;
 use std::sync::Arc;
 use types::{BeaconBlock, BeaconState, BeaconStateError, EthSpec, Hash256, Slot};
 
-/// Extends `BlockRootsIterator`, returning `BeaconBlock` instances, instead of their roots.
+#[derive(Clone)]
-pub struct BlockIterator<T: EthSpec, U> {
+pub struct StateRootsIterator<'a, T: EthSpec, U> {
-    roots: BlockRootsIterator<T, U>,
+    store: Arc<U>,
+    beacon_state: Cow<'a, BeaconState<T>>,
+    slot: Slot,
 }
 
-impl<T: EthSpec, U: Store> BlockIterator<T, U> {
+impl<'a, T: EthSpec, U: Store> StateRootsIterator<'a, 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 {
+    pub fn new(store: Arc<U>, beacon_state: &'a BeaconState<T>, start_slot: Slot) -> Self {
+        Self {
+            store,
+            beacon_state: Cow::Borrowed(beacon_state),
+            slot: start_slot,
+        }
+    }
+}
+
+impl<'a, T: EthSpec, U: Store> Iterator for StateRootsIterator<'a, T, U> {
+    type Item = (Hash256, Slot);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if (self.slot == 0) || (self.slot > self.beacon_state.slot) {
+            return None;
+        }
+
+        self.slot -= 1;
+
+        match self.beacon_state.get_state_root(self.slot) {
+            Ok(root) => Some((*root, self.slot)),
+            Err(BeaconStateError::SlotOutOfBounds) => {
+                // Read a `BeaconState` from the store that has access to prior historical root.
+                let beacon_state: BeaconState<T> = {
+                    let new_state_root = self.beacon_state.get_oldest_state_root().ok()?;
+
+                    self.store.get(&new_state_root).ok()?
+                }?;
+
+                self.beacon_state = Cow::Owned(beacon_state);
+
+                let root = self.beacon_state.get_state_root(self.slot).ok()?;
+
+                Some((*root, self.slot))
+            }
+            _ => None,
+        }
+    }
+}
+
+#[derive(Clone)]
+/// Extends `BlockRootsIterator`, returning `BeaconBlock` instances, instead of their roots.
+pub struct BlockIterator<'a, T: EthSpec, U> {
+    roots: BlockRootsIterator<'a, T, U>,
+}
+
+impl<'a, T: EthSpec, U: Store> BlockIterator<'a, 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: &'a 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> {
+impl<'a, T: EthSpec, U: Store> Iterator for BlockIterator<'a, T, U> {
     type Item = BeaconBlock;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -32,24 +83,25 @@ impl<T: EthSpec, U: Store> Iterator for BlockIterator<T, U> {
 /// exhausted.
 ///
 /// Returns `None` for roots prior to genesis or when there is an error reading from `Store`.
-pub struct BlockRootsIterator<T: EthSpec, U> {
+#[derive(Clone)]
+pub struct BlockRootsIterator<'a, T: EthSpec, U> {
     store: Arc<U>,
-    beacon_state: BeaconState<T>,
+    beacon_state: Cow<'a, BeaconState<T>>,
     slot: Slot,
 }
 
-impl<T: EthSpec, U: Store> BlockRootsIterator<T, U> {
+impl<'a, T: EthSpec, U: Store> BlockRootsIterator<'a, T, U> {
     /// Create a new iterator over all block roots in the given `beacon_state` and prior states.
-    pub fn new(store: Arc<U>, beacon_state: BeaconState<T>, start_slot: Slot) -> Self {
+    pub fn new(store: Arc<U>, beacon_state: &'a BeaconState<T>, start_slot: Slot) -> Self {
         Self {
             slot: start_slot,
-            beacon_state,
+            beacon_state: Cow::Borrowed(beacon_state),
             store,
         }
     }
 }
 
-impl<T: EthSpec, U: Store> Iterator for BlockRootsIterator<T, U> {
+impl<'a, T: EthSpec, U: Store> Iterator for BlockRootsIterator<'a, T, U> {
     type Item = (Hash256, Slot);
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -63,14 +115,16 @@ impl<T: EthSpec, U: Store> Iterator for BlockRootsIterator<T, U> {
             Ok(root) => Some((*root, self.slot)),
             Err(BeaconStateError::SlotOutOfBounds) => {
                 // Read a `BeaconState` from the store that has access to prior historical root.
-                self.beacon_state = {
+                let beacon_state: BeaconState<T> = {
                     // Load the earlier state from disk. Skip forward one slot, because a state
                     // doesn't return it's own state root.
-                    let new_state_root = self.beacon_state.get_state_root(self.slot + 1).ok()?;
+                    let new_state_root = self.beacon_state.get_oldest_state_root().ok()?;
 
                     self.store.get(&new_state_root).ok()?
                 }?;
 
+                self.beacon_state = Cow::Owned(beacon_state);
+
                 let root = self.beacon_state.get_block_root(self.slot).ok()?;
 
                 Some((*root, self.slot))
@@ -97,7 +151,7 @@ mod test {
     }
 
     #[test]
-    fn root_iter() {
+    fn block_root_iter() {
         let store = Arc::new(MemoryStore::open());
         let slots_per_historical_root = MainnetEthSpec::slots_per_historical_root();
 
@@ -120,7 +174,13 @@ mod test {
         state_b.latest_state_roots[0] = state_a_root;
         store.put(&state_a_root, &state_a).unwrap();
 
-        let iter = BlockRootsIterator::new(store.clone(), state_b.clone(), state_b.slot - 1);
+        let iter = BlockRootsIterator::new(store.clone(), &state_b, state_b.slot - 1);
+
+        assert!(
+            iter.clone().find(|(_root, slot)| *slot == 0).is_some(),
+            "iter should contain zero slot"
+        );
+
         let mut collected: Vec<(Hash256, Slot)> = iter.collect();
         collected.reverse();
 
@@ -132,4 +192,68 @@ mod test {
             assert_eq!(collected[i].0, Hash256::from(i as u64));
         }
     }
+
+    #[test]
+    fn state_root_iter() {
+        let store = Arc::new(MemoryStore::open());
+        let slots_per_historical_root = MainnetEthSpec::slots_per_historical_root();
+
+        let mut state_a: BeaconState<MainnetEthSpec> = get_state();
+        let mut state_b: BeaconState<MainnetEthSpec> = get_state();
+
+        state_a.slot = Slot::from(slots_per_historical_root);
+        state_b.slot = Slot::from(slots_per_historical_root * 2);
+
+        let mut hashes = (0..).into_iter().map(|i| Hash256::from(i));
+
+        for slot in 0..slots_per_historical_root {
+            state_a
+                .set_state_root(Slot::from(slot), hashes.next().unwrap())
+                .expect(&format!("should set state_a slot {}", slot));
+        }
+        for slot in slots_per_historical_root..slots_per_historical_root * 2 {
+            state_b
+                .set_state_root(Slot::from(slot), hashes.next().unwrap())
+                .expect(&format!("should set state_b slot {}", slot));
+        }
+
+        /*
+        for root in &mut state_a.latest_state_roots[..] {
+            state_a.set_state_root(slots.next().unwrap(), hashes.next().unwrap());
+            // *root = hashes.next().unwrap()
+        }
+        for root in &mut state_b.latest_state_roots[..] {
+            state_b.set_state_root(slots.next().unwrap(), hashes.next().unwrap());
+            *root = hashes.next().unwrap()
+        }
+        */
+
+        let state_a_root = Hash256::from(slots_per_historical_root as u64);
+        let state_b_root = Hash256::from(slots_per_historical_root as u64 * 2);
+
+        store.put(&state_a_root, &state_a).unwrap();
+        store.put(&state_b_root, &state_b).unwrap();
+
+        let iter = StateRootsIterator::new(store.clone(), &state_b, state_b.slot - 1);
+
+        assert!(
+            iter.clone().find(|(_root, slot)| *slot == 0).is_some(),
+            "iter should contain zero slot"
+        );
+
+        let mut collected: Vec<(Hash256, Slot)> = iter.collect();
+        collected.reverse();
+
+        let expected_len = MainnetEthSpec::slots_per_historical_root() * 2 - 1;
+
+        assert_eq!(collected.len(), expected_len, "collection length incorrect");
+
+        for i in 0..expected_len {
+            let (hash, slot) = collected[i];
+
+            assert_eq!(slot, i as u64, "slot mismatch at {}: {} vs {}", i, slot, i);
+
+            assert_eq!(hash, Hash256::from(i as u64), "hash mismatch at {}", i);
+        }
+    }
 }

eth2/types/src/beacon_state.rs

@@ -569,7 +569,7 @@ impl<T: EthSpec> BeaconState<T> {
     ///
     /// Spec v0.6.3
     fn get_latest_state_roots_index(&self, slot: Slot) -> Result<usize, Error> {
-        if (slot < self.slot) && (self.slot <= slot + self.latest_state_roots.len() as u64) {
+        if (slot < self.slot) && (self.slot <= slot + Slot::from(self.latest_state_roots.len())) {
             Ok(slot.as_usize() % self.latest_state_roots.len())
         } else {
             Err(BeaconStateError::SlotOutOfBounds)
@@ -579,11 +579,23 @@ impl<T: EthSpec> BeaconState<T> {
     /// Gets the state root for some slot.
     ///
     /// Spec v0.6.3
-    pub fn get_state_root(&mut self, slot: Slot) -> Result<&Hash256, Error> {
+    pub fn get_state_root(&self, slot: Slot) -> Result<&Hash256, Error> {
         let i = self.get_latest_state_roots_index(slot)?;
         Ok(&self.latest_state_roots[i])
     }
 
+    /// Gets the oldest (earliest slot) state root.
+    ///
+    /// Spec v0.5.1
+    pub fn get_oldest_state_root(&self) -> Result<&Hash256, Error> {
+        let lookback = std::cmp::min(
+            self.slot - Slot::from(self.latest_state_roots.len()),
+            self.slot,
+        );
+        let i = self.get_latest_state_roots_index(self.slot - lookback)?;
+        Ok(&self.latest_state_roots[i])
+    }
+
     /// Sets the latest state root for slot.
     ///
     /// Spec v0.6.3