Make RootedSlotsIterator for traversing slots on the root fork (#4361)

core/src/blocktree.rs

@@ -32,9 +32,11 @@ use std::sync::mpsc::{sync_channel, Receiver, SyncSender, TrySendError};
 use std::sync::{Arc, RwLock};
 
 pub use self::meta::*;
+pub use self::rooted_slot_iterator::*;
 
 mod db;
 mod meta;
+mod rooted_slot_iterator;
 
 macro_rules! db_imports {
     { $mod:ident, $db:ident, $db_path:expr } => {
@@ -74,6 +76,7 @@ pub enum BlocktreeError {
     RocksDb(rocksdb::Error),
     #[cfg(feature = "kvstore")]
     KvsDb(kvstore::Error),
+    SlotNotRooted,
 }
 
 // ledger window
@@ -182,6 +185,10 @@ impl Blocktree {
         self.orphans_cf.get(slot)
     }
 
+    pub fn rooted_slot_iterator<'a>(&'a self, slot: u64) -> Result<RootedSlotIterator<'a>> {
+        RootedSlotIterator::new(slot, self)
+    }
+
     pub fn slot_meta_iterator(&self, slot: u64) -> Result<impl Iterator<Item = (u64, SlotMeta)>> {
         let meta_iter = self.db.iter::<cf::SlotMeta>(Some(slot))?;
         Ok(meta_iter.map(|(slot, slot_meta_bytes)| {

core/src/blocktree/rooted_slot_iterator.rs

@@ -0,0 +1,126 @@
+use super::*;
+
+pub struct RootedSlotIterator<'a> {
+    next_slots: Vec<u64>,
+    blocktree: &'a super::Blocktree,
+}
+
+impl<'a> RootedSlotIterator<'a> {
+    pub fn new(start_slot: u64, blocktree: &'a super::Blocktree) -> Result<Self> {
+        if blocktree.is_root(start_slot) {
+            Ok(Self {
+                next_slots: vec![start_slot],
+                blocktree,
+            })
+        } else {
+            Err(Error::BlocktreeError(BlocktreeError::SlotNotRooted))
+        }
+    }
+}
+impl<'a> Iterator for RootedSlotIterator<'a> {
+    type Item = (u64, super::SlotMeta);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // Clone b/c passing the closure to the map below requires exclusive access to
+        // `self`, which is borrowed here if we don't clone.
+        let rooted_slot = self
+            .next_slots
+            .iter()
+            .find(|x| self.blocktree.is_root(**x))
+            .cloned();
+
+        rooted_slot.map(|rooted_slot| {
+            let slot_meta = self
+                .blocktree
+                .meta(rooted_slot)
+                .expect("Database failure, couldnt fetch SlotMeta")
+                .expect("SlotMeta in iterator didn't exist");
+
+            self.next_slots = slot_meta.next_slots.clone();
+            (rooted_slot, slot_meta)
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::blocktree_processor::tests::fill_blocktree_slot_with_ticks;
+
+    #[test]
+    fn test_rooted_slot_iterator() {
+        let blocktree_path = get_tmp_ledger_path("test_rooted_slot_iterator");
+        let blocktree = Blocktree::open(&blocktree_path).unwrap();
+        blocktree.set_root(0, 0).unwrap();
+        let ticks_per_slot = 5;
+        /*
+            Build a blocktree in the ledger with the following fork structure:
+
+                 slot 0
+                   |
+                 slot 1  <-- set_root(true)
+                 /   \
+            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_blocktree_slot_with_ticks(
+                &blocktree,
+                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_blocktree_slot_with_ticks(&blocktree, ticks_per_slot, 4, fork_point, fork_hash);
+
+        // Set a root
+        blocktree.set_root(3, 0).unwrap();
+
+        // Trying to get an iterator on a different fork will error
+        assert!(RootedSlotIterator::new(4, &blocktree).is_err());
+
+        // Trying to get an iterator on any slot on the root fork should succeed
+        let result: Vec<_> = RootedSlotIterator::new(3, &blocktree)
+            .unwrap()
+            .into_iter()
+            .map(|(slot, _)| slot)
+            .collect();
+        let expected = vec![3];
+        assert_eq!(result, expected);
+
+        let result: Vec<_> = RootedSlotIterator::new(0, &blocktree)
+            .unwrap()
+            .into_iter()
+            .map(|(slot, _)| slot)
+            .collect();
+        let expected = vec![0, 1, 2, 3];
+        assert_eq!(result, expected);
+
+        drop(blocktree);
+        Blocktree::destroy(&blocktree_path).expect("Expected successful database destruction");
+    }
+}

core/src/blocktree_processor.rs

@@ -281,7 +281,7 @@ pub fn process_blocktree(
 }
 
 #[cfg(test)]
-mod tests {
+pub mod tests {
     use super::*;
     use crate::blocktree::create_new_tmp_ledger;
     use crate::blocktree::tests::entries_to_blobs;
@@ -295,7 +295,7 @@ mod tests {
     use solana_sdk::system_transaction;
     use solana_sdk::transaction::TransactionError;
 
-    fn fill_blocktree_slot_with_ticks(
+    pub fn fill_blocktree_slot_with_ticks(
         blocktree: &Blocktree,
         ticks_per_slot: u64,
         slot: u64,