From f8c425d6b4bb56af4873a1ece9f8a63124bd5aff Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Tue, 21 May 2019 12:58:11 +1000
Subject: [PATCH] Fix and add tests for db get block at slot

---
 beacon_node/db/Cargo.toml           |   3 +-
 beacon_node/db/src/block_at_slot.rs | 137 ++++++++++++++++++++++++++++
 beacon_node/db/src/lib.rs           |   2 +-
 3 files changed, 139 insertions(+), 3 deletions(-)

diff --git a/beacon_node/db/Cargo.toml b/beacon_node/db/Cargo.toml
index bb2f659f8..b6cdafe04 100644
--- a/beacon_node/db/Cargo.toml
+++ b/beacon_node/db/Cargo.toml
@@ -8,10 +8,9 @@ edition = "2018"
 blake2-rfc = "0.2.18"
 bls = { path = "../../eth2/utils/bls" }
 bytes = "0.4.10"
-db_encode = { path = "../db_encode" }
-db_encode_derive = { path = "../db_encode_derive" }
 parking_lot = "0.7"
 rocksdb = "0.10.1"
 ssz = { path = "../../eth2/utils/ssz" }
 ssz_derive = { path = "../../eth2/utils/ssz_derive" }
+tree_hash = { path = "../../eth2/utils/tree_hash" }
 types = { path =  "../../eth2/types" }
diff --git a/beacon_node/db/src/block_at_slot.rs b/beacon_node/db/src/block_at_slot.rs
index c18c8998c..4fa9635a2 100644
--- a/beacon_node/db/src/block_at_slot.rs
+++ b/beacon_node/db/src/block_at_slot.rs
@@ -44,3 +44,140 @@ pub fn get_block_at_preceeding_slot<T: Store>(
         }
     }
 }
+
+#[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 = MemoryDB::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 = MemoryDB::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());
+    }
+}
diff --git a/beacon_node/db/src/lib.rs b/beacon_node/db/src/lib.rs
index 8ac28092a..21b5b0a75 100644
--- a/beacon_node/db/src/lib.rs
+++ b/beacon_node/db/src/lib.rs
@@ -28,8 +28,8 @@ pub trait Store: Sync + Send + Sized {
 
     fn get_block_at_preceeding_slot(
         &self,
-        slot: Slot,
         start_block_root: Hash256,
+        slot: Slot,
     ) -> Result<Option<(Hash256, BeaconBlock)>, Error> {
         block_at_slot::get_block_at_preceeding_slot(self, slot, start_block_root)
     }