From c151e5861e061184fd92c41447afe4559cf2007e Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Tue, 4 Jun 2019 16:39:05 +1000
Subject: [PATCH] epoch processing: fix shard delta calculation

---
 eth2/types/src/beacon_state.rs                    |  7 +++----
 eth2/types/src/beacon_state/beacon_state_types.rs | 10 ++++++++++
 eth2/types/src/beacon_state/committee_cache.rs    | 12 +++++-------
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/eth2/types/src/beacon_state.rs b/eth2/types/src/beacon_state.rs
index 4e510940f0..c4000b8d84 100644
--- a/eth2/types/src/beacon_state.rs
+++ b/eth2/types/src/beacon_state.rs
@@ -300,11 +300,10 @@ impl<T: EthSpec> BeaconState<T> {
 
     pub fn next_epoch_start_shard(&self) -> Result<u64, Error> {
         let cache = self.cache(RelativeEpoch::Current)?;
+        let active_validator_count = cache.active_validator_count();
+        let shard_delta = T::get_shard_delta(active_validator_count);
 
-        Ok(
-            (cache.epoch_start_shard() + cache.epoch_committee_count() as u64)
-                & T::shard_count() as u64,
-        )
+        Ok((self.latest_start_shard + shard_delta) % T::ShardCount::to_u64())
     }
 
     /// Get the slot of an attestation.
diff --git a/eth2/types/src/beacon_state/beacon_state_types.rs b/eth2/types/src/beacon_state/beacon_state_types.rs
index ec6eb68bc4..b8c3a46ed2 100644
--- a/eth2/types/src/beacon_state/beacon_state_types.rs
+++ b/eth2/types/src/beacon_state/beacon_state_types.rs
@@ -29,6 +29,16 @@ pub trait EthSpec: 'static + Default + Sync + Send + Clone + Debug + PartialEq {
         ) * slots_per_epoch
     }
 
+    /// Return the number of shards to increment `state.latest_start_shard` by in a given epoch.
+    ///
+    /// Spec v0.6.3
+    fn get_shard_delta(active_validator_count: usize) -> u64 {
+        std::cmp::min(
+            Self::get_epoch_committee_count(active_validator_count) as u64,
+            Self::ShardCount::to_u64() - Self::ShardCount::to_u64() / Self::spec().slots_per_epoch,
+        )
+    }
+
     /// Returns the minimum number of validators required for this spec.
     ///
     /// This is the _absolute_ minimum, the number required to make the chain operate in the most
diff --git a/eth2/types/src/beacon_state/committee_cache.rs b/eth2/types/src/beacon_state/committee_cache.rs
index 27374c3391..b5de99c835 100644
--- a/eth2/types/src/beacon_state/committee_cache.rs
+++ b/eth2/types/src/beacon_state/committee_cache.rs
@@ -49,19 +49,17 @@ impl CommitteeCache {
         let shuffling_start_shard = match relative_epoch {
             RelativeEpoch::Current => state.latest_start_shard,
             RelativeEpoch::Previous => {
-                let committees_in_previous_epoch =
-                    T::get_epoch_committee_count(active_validator_indices.len()) as u64;
+                let shard_delta = T::get_shard_delta(active_validator_indices.len());
 
-                (state.latest_start_shard + T::shard_count() as u64 - committees_in_previous_epoch)
-                    % T::shard_count() as u64
+                (state.latest_start_shard + T::ShardCount::to_u64() - shard_delta)
+                    % T::ShardCount::to_u64()
             }
             RelativeEpoch::Next => {
                 let current_active_validators =
                     get_active_validator_count(&state.validator_registry, state.current_epoch());
-                let committees_in_current_epoch =
-                    T::get_epoch_committee_count(current_active_validators) as u64;
+                let shard_delta = T::get_shard_delta(current_active_validators);
 
-                (state.latest_start_shard + committees_in_current_epoch) % T::shard_count() as u64
+                (state.latest_start_shard + shard_delta) % T::ShardCount::to_u64()
             }
         };