Add state roots iter to store
This commit is contained in:
		
							parent
							
								
									2b5c70711d
								
							
						
					
					
						commit
						952e08ba38
					
				| @ -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); | ||||
|         } | ||||
|     } | ||||
| } | ||||
|  | ||||
| @ -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
 | ||||
|  | ||||
		Loading…
	
		Reference in New Issue
	
	Block a user