From 952e08ba3898194a07243d7a87d4b8d2571fe2df Mon Sep 17 00:00:00 2001 From: Paul Hauner Date: Wed, 19 Jun 2019 01:47:21 +1000 Subject: [PATCH] Add state roots iter to store --- beacon_node/store/src/iter.rs | 156 +++++++++++++++++++++++++++++---- eth2/types/src/beacon_state.rs | 16 +++- 2 files changed, 154 insertions(+), 18 deletions(-) diff --git a/beacon_node/store/src/iter.rs b/beacon_node/store/src/iter.rs index 844837983b..2326052123 100644 --- a/beacon_node/store/src/iter.rs +++ b/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. -pub struct BlockIterator { - roots: BlockRootsIterator, +#[derive(Clone)] +pub struct StateRootsIterator<'a, T: EthSpec, U> { + store: Arc, + beacon_state: Cow<'a, BeaconState>, + slot: Slot, } -impl BlockIterator { +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, beacon_state: BeaconState, start_slot: Slot) -> Self { + pub fn new(store: Arc, beacon_state: &'a BeaconState, 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 { + 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 = { + 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, beacon_state: &'a BeaconState, start_slot: Slot) -> Self { Self { roots: BlockRootsIterator::new(store, beacon_state, start_slot), } } } -impl Iterator for BlockIterator { +impl<'a, T: EthSpec, U: Store> Iterator for BlockIterator<'a, T, U> { type Item = BeaconBlock; fn next(&mut self) -> Option { @@ -32,24 +83,25 @@ impl Iterator for BlockIterator { /// exhausted. /// /// Returns `None` for roots prior to genesis or when there is an error reading from `Store`. -pub struct BlockRootsIterator { +#[derive(Clone)] +pub struct BlockRootsIterator<'a, T: EthSpec, U> { store: Arc, - beacon_state: BeaconState, + beacon_state: Cow<'a, BeaconState>, slot: Slot, } -impl BlockRootsIterator { +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, beacon_state: BeaconState, start_slot: Slot) -> Self { + pub fn new(store: Arc, beacon_state: &'a BeaconState, start_slot: Slot) -> Self { Self { slot: start_slot, - beacon_state, + beacon_state: Cow::Borrowed(beacon_state), store, } } } -impl Iterator for BlockRootsIterator { +impl<'a, T: EthSpec, U: Store> Iterator for BlockRootsIterator<'a, T, U> { type Item = (Hash256, Slot); fn next(&mut self) -> Option { @@ -63,14 +115,16 @@ impl Iterator for BlockRootsIterator { 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 = { // 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 = get_state(); + let mut state_b: BeaconState = 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); + } + } } diff --git a/eth2/types/src/beacon_state.rs b/eth2/types/src/beacon_state.rs index f76676a2b5..58237fe6f3 100644 --- a/eth2/types/src/beacon_state.rs +++ b/eth2/types/src/beacon_state.rs @@ -569,7 +569,7 @@ impl BeaconState { /// /// Spec v0.6.3 fn get_latest_state_roots_index(&self, slot: Slot) -> Result { - 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 BeaconState { /// 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