Overhaul state diffing

consensus/state_processing/src/consensus_context.rs

@@ -6,8 +6,8 @@ use std::collections::{hash_map::Entry, HashMap};
 use std::marker::PhantomData;
 use tree_hash::TreeHash;
 use types::{
-    Attestation, AttestationData, BeaconState, BeaconStateError, BitList, ChainSpec, EthSpec,
-    ExecPayload, Hash256, IndexedAttestation, SignedBeaconBlock, Slot,
+    Attestation, AttestationData, BeaconState, BeaconStateError, BitList, ChainSpec, Epoch,
+    EthSpec, ExecPayload, Hash256, IndexedAttestation, SignedBeaconBlock, Slot,
 };
 
 #[derive(Debug, Clone)]
@@ -31,6 +31,7 @@ pub enum ContextError {
     BeaconState(BeaconStateError),
     EpochCache(EpochCacheError),
     SlotMismatch { slot: Slot, expected: Slot },
+    EpochMismatch { epoch: Epoch, expected: Epoch },
 }
 
 impl From<BeaconStateError> for ContextError {
@@ -62,12 +63,13 @@ impl<T: EthSpec> ConsensusContext<T> {
         self
     }
 
+    // FIXME(sproul): extra safety checks?
     pub fn get_proposer_index(
         &mut self,
         state: &BeaconState<T>,
         spec: &ChainSpec,
     ) -> Result<u64, ContextError> {
-        self.check_slot(state.slot())?;
+        self.check_epoch(state.current_epoch())?;
 
         if let Some(proposer_index) = self.proposer_index {
             return Ok(proposer_index);
@@ -109,6 +111,15 @@ impl<T: EthSpec> ConsensusContext<T> {
         }
     }
 
+    fn check_epoch(&self, epoch: Epoch) -> Result<(), ContextError> {
+        let expected = self.slot.epoch(T::slots_per_epoch());
+        if epoch == expected {
+            Ok(())
+        } else {
+            Err(ContextError::EpochMismatch { epoch, expected })
+        }
+    }
+
     pub fn set_epoch_cache(mut self, epoch_cache: EpochCache) -> Self {
         self.epoch_cache = Some(epoch_cache);
         self