Add state roots iter to store

This commit is contained in:
Paul Hauner 2019-06-19 01:47:21 +10:00
parent 2b5c70711d
commit 952e08ba38
No known key found for this signature in database
GPG Key ID: 303E4494BB28068C
2 changed files with 154 additions and 18 deletions

View File

@ -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<T: EthSpec, U> {
roots: BlockRootsIterator<T, U>,
#[derive(Clone)]
pub struct StateRootsIterator<'a, T: EthSpec, 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);
}
}
}

View File

@ -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