From 94a39c3cb6c0a297c4c20bbacf18bf70a3954f35 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Sun, 19 May 2019 17:21:24 +1000
Subject: [PATCH] Implement state.get_crosslink_committees_at_slot

---
 eth2/types/src/beacon_state.rs             | 14 ++++++---
 eth2/types/src/beacon_state/epoch_cache.rs | 36 +++++++++++++++++++---
 2 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/eth2/types/src/beacon_state.rs b/eth2/types/src/beacon_state.rs
index 66375a1d2..50f21092d 100644
--- a/eth2/types/src/beacon_state.rs
+++ b/eth2/types/src/beacon_state.rs
@@ -44,6 +44,7 @@ pub enum Error {
     InsufficientSlashedBalances,
     InsufficientStateRoots,
     NoCommitteeForShard,
+    NoCommitteeForSlot,
     PubkeyCacheInconsistent,
     PubkeyCacheIncomplete {
         cache_len: usize,
@@ -346,10 +347,15 @@ impl<T: EthSpec> BeaconState<T> {
     /// Spec v0.5.1
     pub fn get_crosslink_committees_at_slot(
         &self,
-        _slot: Slot,
-        _spec: &ChainSpec,
-    ) -> Result<&Vec<CrosslinkCommittee>, Error> {
-        unimplemented!("FIXME(sproul)")
+        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)?;
+
+        cache
+            .get_crosslink_committees_for_slot(slot)
+            .ok_or_else(|| Error::NoCommitteeForSlot)
     }
 
     /// Returns the crosslink committees for some shard in some cached epoch.
diff --git a/eth2/types/src/beacon_state/epoch_cache.rs b/eth2/types/src/beacon_state/epoch_cache.rs
index 5072703a6..811387d2f 100644
--- a/eth2/types/src/beacon_state/epoch_cache.rs
+++ b/eth2/types/src/beacon_state/epoch_cache.rs
@@ -150,18 +150,44 @@ impl EpochCache {
         self.shuffling_start_shard
     }
 
+    pub fn get_crosslink_committees_for_slot(&self, slot: Slot) -> Option<Vec<CrosslinkCommittee>> {
+        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 {
+            let mut committees = Vec::with_capacity(committees_per_slot);
+
+            for index in position..position + committees_per_slot {
+                let committee = self.compute_committee(index)?;
+                let shard = (self.shuffling_start_shard + index as u64) % self.shard_count;
+
+                committees.push(CrosslinkCommittee {
+                    committee,
+                    shard,
+                    slot,
+                });
+            }
+
+            Some(committees)
+        }
+    }
+
     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 {
+        if position >= self.committee_count {
             None
         } else {
-            let committees_per_slot = self.committee_count / self.slots_per_epoch as usize;
-            let index = position * committees_per_slot;
-
-            self.compute_committee(index)
+            self.compute_committee(position)
         }
     }