Tidy epoch processing, improve testing

eth2/types/src/beacon_state.rs

@@ -348,7 +348,6 @@ impl<T: EthSpec> BeaconState<T> {
     pub fn get_crosslink_committees_at_slot(
         &self,
         slot: Slot,
-        spec: &ChainSpec,
     ) -> Result<Vec<CrosslinkCommittee>, Error> {
         let relative_epoch = RelativeEpoch::from_slot(self.slot, slot, T::slots_per_epoch())?;
         let cache = self.cache(relative_epoch)?;

eth2/types/src/beacon_state/epoch_cache.rs

@@ -178,17 +178,9 @@ impl EpochCache {
     }
 
     pub fn first_committee_at_slot(&self, slot: Slot) -> Option<&[usize]> {
-        let position = self
-            .initialized_epoch?
-            .position(slot, self.slots_per_epoch)?;
-        let committees_per_slot = self.committee_count / self.slots_per_epoch as usize;
-        let position = position * committees_per_slot;
-
-        if position >= self.committee_count {
-            None
-        } else {
-            self.compute_committee(position)
-        }
+        self.get_crosslink_committees_for_slot(slot)?
+            .first()
+            .and_then(|cc| Some(cc.committee))
     }
 
     fn compute_committee(&self, index: usize) -> Option<&[usize]> {

eth2/types/src/beacon_state/epoch_cache/tests.rs

@@ -30,16 +30,34 @@ fn execute_sane_cache_test<T: EthSpec>(
         shuffle_list(active_indices, spec.shuffle_round_count, &seed[..], false).unwrap();
 
     let mut expected_indices_iter = shuffling.iter();
+    let mut expected_shards_iter = (start_shard..start_shard + T::shard_count() as u64).into_iter();
 
-    for i in 0..T::shard_count() {
-        let shard = (i + start_shard as usize) % T::shard_count();
+    // Loop through all slots in the epoch being tested.
+    for slot in epoch.slot_iter(spec.slots_per_epoch) {
+        let crosslink_committees = state.get_crosslink_committees_at_slot(slot).unwrap();
 
-        let c = state
-            .get_crosslink_committee_for_shard(shard as u64, relative_epoch)
-            .unwrap()
-            .unwrap();
+        // Assert that the number of committees in this slot is consistent with the reported number
+        // of committees in an epoch.
+        assert_eq!(
+            crosslink_committees.len() as u64,
+            state.get_epoch_committee_count(relative_epoch).unwrap() / T::slots_per_epoch()
+        );
 
-        for &i in c.committee {
+        for cc in crosslink_committees {
+            // Assert that shards are assigned contiguously across committees.
+            assert_eq!(expected_shards_iter.next().unwrap(), cc.shard);
+            // Assert that a committee lookup via slot is identical to a committee lookup via
+            // shard.
+            assert_eq!(
+                state
+                    .get_crosslink_committee_for_shard(cc.shard, relative_epoch)
+                    .unwrap(),
+                cc
+            );
+
+            // Loop through each validator in the committee.
+            for &i in cc.committee {
+                // Assert the validators are assigned contiguously across committees.
                 assert_eq!(
                     i,
                     *expected_indices_iter.next().unwrap(),
@@ -49,6 +67,13 @@ fn execute_sane_cache_test<T: EthSpec>(
         }
     }
 
+    // Assert that all validators were assigned to a committee.
+    assert!(expected_indices_iter.next().is_none());
+
+    // Assert that all shards were assigned to a committee.
+    assert!(expected_shards_iter.next().is_none());
+}
+
 fn sane_cache_test<T: EthSpec>(
     validator_count: usize,
     state_epoch: Epoch,

eth2/types/src/test_utils/builders/testing_beacon_block_builder.rs

@@ -109,7 +109,7 @@ impl TestingBeaconBlockBuilder {
                 break;
             }
 
-            for crosslink_committee in state.get_crosslink_committees_at_slot(slot, spec)? {
+            for crosslink_committee in state.get_crosslink_committees_at_slot(slot)? {
                 if attestations_added >= num_attestations {
                     break;
                 }

eth2/types/src/test_utils/builders/testing_beacon_state_builder.rs

@@ -223,7 +223,7 @@ impl<T: EthSpec> TestingBeaconStateBuilder<T> {
             let slot = Slot::from(slot);
 
             let committees: Vec<OwnedCrosslinkCommittee> = state
-                .get_crosslink_committees_at_slot(slot, spec)
+                .get_crosslink_committees_at_slot(slot)
                 .unwrap()
                 .into_iter()
                 .map(|c| c.clone().into_owned())