diff --git a/ledger/src/lib.rs b/ledger/src/lib.rs index dde02ecdb8..af450cc341 100644 --- a/ledger/src/lib.rs +++ b/ledger/src/lib.rs @@ -12,6 +12,7 @@ pub mod genesis_utils; pub mod leader_schedule; pub mod leader_schedule_cache; pub mod leader_schedule_utils; +pub mod next_slots_iterator; pub mod poh; pub mod rooted_slot_iterator; pub mod shred; diff --git a/ledger/src/next_slots_iterator.rs b/ledger/src/next_slots_iterator.rs new file mode 100644 index 0000000000..530ce53454 --- /dev/null +++ b/ledger/src/next_slots_iterator.rs @@ -0,0 +1,117 @@ +use crate::{blockstore::*, blockstore_meta::SlotMeta}; +use solana_sdk::clock::Slot; + +pub struct NextSlotsIterator<'a> { + pending_slots: Vec, + blockstore: &'a Blockstore, +} + +impl<'a> NextSlotsIterator<'a> { + pub fn new(start_slot: Slot, blockstore: &'a Blockstore) -> Self { + Self { + pending_slots: vec![start_slot], + blockstore, + } + } +} + +impl<'a> Iterator for NextSlotsIterator<'a> { + type Item = (Slot, SlotMeta); + fn next(&mut self) -> Option { + if self.pending_slots.is_empty() { + None + } else { + let slot = self.pending_slots.pop().unwrap(); + if let Some(slot_meta) = self.blockstore.meta(slot).unwrap() { + self.pending_slots.extend(slot_meta.next_slots.iter()); + Some((slot, slot_meta)) + } else { + None + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::blockstore_processor::fill_blockstore_slot_with_ticks; + use solana_sdk::hash::Hash; + use std::collections::HashSet; + + #[test] + fn test_next_slots_iterator() { + let blockstore_path = get_tmp_ledger_path!(); + let blockstore = Blockstore::open(&blockstore_path).unwrap(); + blockstore.set_roots(&[0]).unwrap(); + let ticks_per_slot = 5; + /* + Build a blockstore in the ledger with the following fork structure: + + slot 0 + | + slot 1 <-- set_root + / \ + slot 2 | + / | + slot 3 | + | + slot 4 + + */ + + // Fork 1, ending at slot 3 + let last_entry_hash = Hash::default(); + let fork_point = 1; + let mut fork_hash = Hash::default(); + for slot in 0..=3 { + let parent = { + if slot == 0 { + 0 + } else { + slot - 1 + } + }; + let last_entry_hash = fill_blockstore_slot_with_ticks( + &blockstore, + ticks_per_slot, + slot, + parent, + last_entry_hash, + ); + + if slot == fork_point { + fork_hash = last_entry_hash; + } + } + + // Fork 2, ending at slot 4 + let _ = + fill_blockstore_slot_with_ticks(&blockstore, ticks_per_slot, 4, fork_point, fork_hash); + + // Trying to get an iterator on any slot on the root fork should succeed + let result: HashSet<_> = NextSlotsIterator::new(0, &blockstore) + .into_iter() + .map(|(slot, _)| slot) + .collect(); + let expected = vec![0, 1, 2, 3, 4].into_iter().collect(); + assert_eq!(result, expected); + + let result: HashSet<_> = NextSlotsIterator::new(2, &blockstore) + .into_iter() + .map(|(slot, _)| slot) + .collect(); + let expected = vec![2, 3].into_iter().collect(); + assert_eq!(result, expected); + + let result: HashSet<_> = NextSlotsIterator::new(4, &blockstore) + .into_iter() + .map(|(slot, _)| slot) + .collect(); + let expected = vec![4].into_iter().collect(); + assert_eq!(result, expected); + + drop(blockstore); + Blockstore::destroy(&blockstore_path).expect("Expected successful database destruction"); + } +}