Rename db crate to store

beacon_node/store/src/block_at_slot.rs

@@ -0,0 +1,183 @@
+use super::*;
+use ssz::{Decode, DecodeError};
+
+fn get_block_bytes<T: Store>(store: &T, root: Hash256) -> Result<Option<Vec<u8>>, Error> {
+    store.get_bytes(BeaconBlock::db_column().into(), &root[..])
+}
+
+fn read_slot_from_block_bytes(bytes: &[u8]) -> Result<Slot, DecodeError> {
+    let end = std::cmp::min(Slot::ssz_fixed_len(), bytes.len());
+
+    Slot::from_ssz_bytes(&bytes[0..end])
+}
+
+fn read_previous_block_root_from_block_bytes(bytes: &[u8]) -> Result<Hash256, DecodeError> {
+    let previous_bytes = Slot::ssz_fixed_len();
+    let slice = bytes
+        .get(previous_bytes..previous_bytes + Hash256::ssz_fixed_len())
+        .ok_or_else(|| DecodeError::BytesInvalid("Not enough bytes.".to_string()))?;
+
+    Hash256::from_ssz_bytes(slice)
+}
+
+pub fn get_block_at_preceeding_slot<T: Store>(
+    store: &T,
+    slot: Slot,
+    start_root: Hash256,
+) -> Result<Option<(Hash256, BeaconBlock)>, Error> {
+    let mut root = start_root;
+
+    loop {
+        if let Some(bytes) = get_block_bytes(store, root)? {
+            let this_slot = read_slot_from_block_bytes(&bytes)?;
+
+            if this_slot == slot {
+                let block = BeaconBlock::from_ssz_bytes(&bytes)?;
+                break Ok(Some((root, block)));
+            } else if this_slot < slot {
+                break Ok(None);
+            } else {
+                root = read_previous_block_root_from_block_bytes(&bytes)?;
+            }
+        } else {
+            break Ok(None);
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use ssz::Encode;
+    use tree_hash::TreeHash;
+
+    #[test]
+    fn read_slot() {
+        let spec = FewValidatorsEthSpec::spec();
+
+        let test_slot = |slot: Slot| {
+            let mut block = BeaconBlock::empty(&spec);
+            block.slot = slot;
+            let bytes = block.as_ssz_bytes();
+            assert_eq!(read_slot_from_block_bytes(&bytes).unwrap(), slot);
+        };
+
+        test_slot(Slot::new(0));
+        test_slot(Slot::new(1));
+        test_slot(Slot::new(42));
+        test_slot(Slot::new(u64::max_value()));
+    }
+
+    #[test]
+    fn bad_slot() {
+        for i in 0..8 {
+            assert!(read_slot_from_block_bytes(&vec![0; i]).is_err());
+        }
+    }
+
+    #[test]
+    fn read_previous_block_root() {
+        let spec = FewValidatorsEthSpec::spec();
+
+        let test_root = |root: Hash256| {
+            let mut block = BeaconBlock::empty(&spec);
+            block.previous_block_root = root;
+            let bytes = block.as_ssz_bytes();
+            assert_eq!(
+                read_previous_block_root_from_block_bytes(&bytes).unwrap(),
+                root
+            );
+        };
+
+        test_root(Hash256::random());
+        test_root(Hash256::random());
+        test_root(Hash256::random());
+    }
+
+    fn build_chain(
+        store: &impl Store,
+        slots: &[usize],
+        spec: &ChainSpec,
+    ) -> Vec<(Hash256, BeaconBlock)> {
+        let mut blocks_and_roots: Vec<(Hash256, BeaconBlock)> = vec![];
+
+        for (i, slot) in slots.iter().enumerate() {
+            let mut block = BeaconBlock::empty(spec);
+            block.slot = Slot::from(*slot);
+
+            if i > 0 {
+                block.previous_block_root = blocks_and_roots[i - 1].0;
+            }
+
+            let root = Hash256::from_slice(&block.tree_hash_root());
+
+            store.put(&root, &block).unwrap();
+            blocks_and_roots.push((root, block));
+        }
+
+        blocks_and_roots
+    }
+
+    #[test]
+    fn chain_without_skips() {
+        let n: usize = 10;
+        let store = MemoryStore::open();
+        let spec = FewValidatorsEthSpec::spec();
+
+        let slots: Vec<usize> = (0..n).collect();
+        let blocks_and_roots = build_chain(&store, &slots, &spec);
+
+        for source in 1..n {
+            for target in 0..=source {
+                let (source_root, _source_block) = &blocks_and_roots[source];
+                let (target_root, target_block) = &blocks_and_roots[target];
+
+                let (found_root, found_block) = store
+                    .get_block_at_preceeding_slot(*source_root, target_block.slot)
+                    .unwrap()
+                    .unwrap();
+
+                assert_eq!(found_root, *target_root);
+                assert_eq!(found_block, *target_block);
+            }
+        }
+    }
+
+    #[test]
+    fn chain_with_skips() {
+        let store = MemoryStore::open();
+        let spec = FewValidatorsEthSpec::spec();
+
+        let slots = vec![0, 1, 2, 5];
+
+        let blocks_and_roots = build_chain(&store, &slots, &spec);
+
+        // Valid slots
+        for target in 0..3 {
+            let (source_root, _source_block) = &blocks_and_roots[3];
+            let (target_root, target_block) = &blocks_and_roots[target];
+
+            let (found_root, found_block) = store
+                .get_block_at_preceeding_slot(*source_root, target_block.slot)
+                .unwrap()
+                .unwrap();
+
+            assert_eq!(found_root, *target_root);
+            assert_eq!(found_block, *target_block);
+        }
+
+        // Slot that doesn't exist
+        let (source_root, _source_block) = &blocks_and_roots[3];
+        assert!(store
+            .get_block_at_preceeding_slot(*source_root, Slot::new(3))
+            .unwrap()
+            .is_none());
+
+        // Slot too high
+        let (source_root, _source_block) = &blocks_and_roots[3];
+        assert!(store
+            .get_block_at_preceeding_slot(*source_root, Slot::new(3))
+            .unwrap()
+            .is_none());
+    }
+}