From 29a3e0c868456c32583c5c9588996be303cd3bed Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Sun, 19 May 2019 16:03:54 +1000
Subject: [PATCH] Move beacon state fns to be `RelativeEpoch` native

---
 eth2/types/src/beacon_state.rs                | 37 +++++++++++--------
 .../src/beacon_state/epoch_cache/tests.rs     |  5 ++-
 2 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/eth2/types/src/beacon_state.rs b/eth2/types/src/beacon_state.rs
index 756e02d72..42d88242d 100644
--- a/eth2/types/src/beacon_state.rs
+++ b/eth2/types/src/beacon_state.rs
@@ -50,6 +50,7 @@ pub enum Error {
     },
     PreviousEpochCacheUninitialized,
     CurrentEpochCacheUnintialized,
+    RelativeEpochError(RelativeEpochError),
     EpochCacheUnintialized(RelativeEpoch),
     EpochCacheError(EpochCacheError),
     TreeHashCacheError(TreeHashCacheError),
@@ -284,11 +285,11 @@ impl<T: EthSpec> BeaconState<T> {
     ///
     /// Spec v0.6.1
     pub fn get_attestation_slot(&self, attestation_data: &AttestationData) -> Result<Slot, Error> {
+        let target_epoch =
+            RelativeEpoch::from_epoch(self.current_epoch(), attestation_data.target_epoch)?;
+
         let cc = self
-            .get_crosslink_committee_for_shard(
-                attestation_data.shard,
-                attestation_data.target_epoch,
-            )?
+            .get_crosslink_committee_for_shard(attestation_data.shard, target_epoch)?
             .ok_or_else(|| Error::NoCommitteeForShard)?;
         Ok(cc.slot)
     }
@@ -298,8 +299,11 @@ impl<T: EthSpec> BeaconState<T> {
     /// Note: the indices are shuffled (i.e., not in ascending order).
     ///
     /// Returns an error if that epoch is not cached, or the cache is not initialized.
-    pub fn get_cached_active_validator_indices(&self, epoch: Epoch) -> Result<&[usize], Error> {
-        let cache = self.cache(epoch)?;
+    pub fn get_cached_active_validator_indices(
+        &self,
+        relative_epoch: RelativeEpoch,
+    ) -> Result<&[usize], Error> {
+        let cache = self.cache(relative_epoch)?;
 
         Ok(&cache.active_validator_indices())
     }
@@ -334,9 +338,9 @@ impl<T: EthSpec> BeaconState<T> {
     pub fn get_crosslink_committee_for_shard(
         &self,
         shard: u64,
-        epoch: Epoch,
+        relative_epoch: RelativeEpoch,
     ) -> Result<Option<CrosslinkCommittee>, Error> {
-        let cache = self.cache(epoch)?;
+        let cache = self.cache(relative_epoch)?;
 
         Ok(cache.get_crosslink_committee_for_shard(shard))
     }
@@ -662,7 +666,7 @@ impl<T: EthSpec> BeaconState<T> {
     pub fn get_churn_limit(&self, spec: &ChainSpec) -> Result<u64, Error> {
         Ok(std::cmp::max(
             spec.min_per_epoch_churn_limit,
-            self.cache(self.current_epoch())?.active_validator_count() as u64
+            self.cache(RelativeEpoch::Current)?.active_validator_count() as u64
                 / spec.churn_limit_quotient,
         ))
     }
@@ -679,7 +683,7 @@ impl<T: EthSpec> BeaconState<T> {
         &self,
         validator_index: usize,
     ) -> Result<&Option<AttestationDuty>, Error> {
-        let cache = self.cache(self.current_epoch())?;
+        let cache = self.cache(RelativeEpoch::Current)?;
 
         Ok(cache
             .attestation_duties
@@ -765,13 +769,10 @@ impl<T: EthSpec> BeaconState<T> {
 
     /// Returns the cache for some `RelativeEpoch`. Returns an error if the cache has not been
     /// initialized.
-    fn cache(&self, epoch: Epoch) -> Result<&EpochCache, Error> {
-        let relative_epoch = RelativeEpoch::from_epoch(self.current_epoch(), epoch)
-            .map_err(|_| Error::EpochOutOfBounds)?;
-
+    fn cache(&self, relative_epoch: RelativeEpoch) -> Result<&EpochCache, Error> {
         let cache = &self.epoch_caches[Self::cache_index(relative_epoch)];
 
-        if cache.is_initialized_at(epoch) {
+        if cache.is_initialized_at(relative_epoch.into_epoch(self.current_epoch())) {
             Ok(cache)
         } else {
             Err(Error::EpochCacheUnintialized(relative_epoch))
@@ -845,6 +846,12 @@ impl From<EpochCacheError> for Error {
     }
 }
 
+impl From<RelativeEpochError> for Error {
+    fn from(e: RelativeEpochError) -> Error {
+        Error::RelativeEpochError(e)
+    }
+}
+
 impl From<TreeHashCacheError> for Error {
     fn from(e: TreeHashCacheError) -> Error {
         Error::TreeHashCacheError(e)
diff --git a/eth2/types/src/beacon_state/epoch_cache/tests.rs b/eth2/types/src/beacon_state/epoch_cache/tests.rs
index 7b218dd86..31b84b57a 100644
--- a/eth2/types/src/beacon_state/epoch_cache/tests.rs
+++ b/eth2/types/src/beacon_state/epoch_cache/tests.rs
@@ -14,9 +14,10 @@ fn execute_sane_cache_test<T: EthSpec>(
     let active_indices: Vec<usize> = (0..validator_count).collect();
     let seed = state.generate_seed(epoch, spec).unwrap();
     let start_shard = 0;
+    let relative_epoch = RelativeEpoch::from_epoch(state.current_epoch(), epoch).unwrap();
 
     let mut ordered_indices = state
-        .get_cached_active_validator_indices(epoch)
+        .get_cached_active_validator_indices(relative_epoch)
         .unwrap()
         .to_vec();
     ordered_indices.sort_unstable();
@@ -34,7 +35,7 @@ fn execute_sane_cache_test<T: EthSpec>(
         let shard = (i + start_shard as usize) % T::shard_count();
 
         let c = state
-            .get_crosslink_committee_for_shard(shard as u64, epoch)
+            .get_crosslink_committee_for_shard(shard as u64, relative_epoch)
             .unwrap()
             .unwrap();