Update and add tests for chain maps

beacon_chain/chain/src/lib.rs

@@ -78,7 +78,7 @@ impl<T> BeaconChain<T>
         let mut attester_proposer_maps = HashMap::new();
 
         let (attester_map, proposer_map) = generate_attester_and_proposer_maps(
-            &crystallized_state, 0)?;
+            &crystallized_state.shard_and_committee_for_slots, 0)?;
 
         active_states.insert(canonical_latest_block_hash, active_state);
         crystallized_states.insert(canonical_latest_block_hash, crystallized_state);

beacon_chain/chain/src/maps.rs

@@ -1,7 +1,7 @@
 use types::{
     AttesterMap,
-    CrystallizedState,
     ProposerMap,
+    ShardAndCommittee,
 };
 
 #[derive(Debug, PartialEq)]
@@ -13,25 +13,24 @@ pub enum AttesterAndProposerMapError {
 /// Generate a map of `(slot, shard) |--> committee`.
 ///
 /// The attester map is used to optimise the lookup of a committee.
-pub fn generate_attester_and_proposer_maps(cry_state: &CrystallizedState, start_slot: u64)
+pub fn generate_attester_and_proposer_maps(
+    shard_and_committee_for_slots: &Vec<Vec<ShardAndCommittee>>,
+    start_slot: u64)
     -> Result<(AttesterMap, ProposerMap), AttesterAndProposerMapError>
 {
     let mut attester_map = AttesterMap::new();
     let mut proposer_map = ProposerMap::new();
-    for (i, slot) in cry_state.shard_and_committee_for_slots.iter().enumerate() {
+    for (i, slot) in shard_and_committee_for_slots.iter().enumerate() {
         /*
          * Store the proposer for the block.
          */
         let slot_number = (i as u64).saturating_add(start_slot);
-        let first_committee = {
-            let first_shard_and_committee = slot.get(0)
-                .ok_or(AttesterAndProposerMapError::NoShardAndCommitteeForSlot)?;
-            first_shard_and_committee.committee.clone()
-        };
-        println!("{:?}", slot);
+        let first_committee = &slot.get(0)
+            .ok_or(AttesterAndProposerMapError::NoShardAndCommitteeForSlot)?
+            .committee;
         let proposer_index = (slot_number as usize).checked_rem(first_committee.len())
             .ok_or(AttesterAndProposerMapError::NoAvailableProposer)?;
-        proposer_map.insert(slot_number, proposer_index);
+        proposer_map.insert(slot_number, first_committee[proposer_index]);
 
         /*
          * Loop through the shards and extend the attester map.
@@ -43,3 +42,83 @@ pub fn generate_attester_and_proposer_maps(cry_state: &CrystallizedState, start_
     };
     Ok((attester_map, proposer_map))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn sac_generator(shard_count: u16,
+                     slot_count: usize,
+                     sac_per_slot: usize,
+                     committee_size: usize)
+        -> Vec<Vec<ShardAndCommittee>>
+    {
+        let mut shard = 0;
+        let mut validator = 0;
+        let mut cycle = vec![];
+
+        for _ in 0..slot_count {
+            let mut slot: Vec<ShardAndCommittee> = vec![];
+            for _ in 0..sac_per_slot {
+                let mut sac = ShardAndCommittee {
+                    shard: shard % shard_count,
+                    committee: vec![],
+                };
+                for _ in 0..committee_size {
+                    sac.committee.push(validator);
+                    validator += 1;
+                }
+                slot.push(sac);
+                shard += 1;
+            }
+            cycle.push(slot);
+        }
+        cycle
+    }
+
+    #[test]
+    fn test_attester_proposer_maps_empty_slots() {
+        let sac = sac_generator(4, 4, 0, 1);
+        let result = generate_attester_and_proposer_maps(&sac, 0);
+        assert_eq!(result, Err(AttesterAndProposerMapError::NoShardAndCommitteeForSlot));
+    }
+
+    #[test]
+    fn test_attester_proposer_maps_empty_committees() {
+        let sac = sac_generator(4, 4, 1, 0);
+        let result = generate_attester_and_proposer_maps(&sac, 0);
+        assert_eq!(result, Err(AttesterAndProposerMapError::NoAvailableProposer));
+    }
+
+    #[test]
+    fn test_attester_proposer_maps_scenario_a() {
+        let sac = sac_generator(4, 4, 1, 1);
+        let (a, p) = generate_attester_and_proposer_maps(&sac, 0).unwrap();
+
+        assert_eq!(*p.get(&0).unwrap(), 0);
+        assert_eq!(*p.get(&1).unwrap(), 1);
+        assert_eq!(*p.get(&2).unwrap(), 2);
+        assert_eq!(*p.get(&3).unwrap(), 3);
+
+        assert_eq!(*a.get(&(0, 0)).unwrap(), vec![0]);
+        assert_eq!(*a.get(&(1, 1)).unwrap(), vec![1]);
+        assert_eq!(*a.get(&(2, 2)).unwrap(), vec![2]);
+        assert_eq!(*a.get(&(3, 3)).unwrap(), vec![3]);
+    }
+
+    #[test]
+    fn test_attester_proposer_maps_scenario_b() {
+        let sac = sac_generator(4, 4, 1, 4);
+        let (a, p) = generate_attester_and_proposer_maps(&sac, 0).unwrap();
+
+        assert_eq!(*p.get(&0).unwrap(), 0);
+        assert_eq!(*p.get(&1).unwrap(), 5);
+        assert_eq!(*p.get(&2).unwrap(), 10);
+        assert_eq!(*p.get(&3).unwrap(), 15);
+
+        assert_eq!(*a.get(&(0, 0)).unwrap(), vec![0, 1, 2, 3]);
+        assert_eq!(*a.get(&(1, 1)).unwrap(), vec![4, 5, 6, 7]);
+        assert_eq!(*a.get(&(2, 2)).unwrap(), vec![8, 9, 10, 11]);
+        assert_eq!(*a.get(&(3, 3)).unwrap(), vec![12, 13, 14, 15]);
+    }
+}