Merged with unstable

consensus/state_processing/src/consensus_context.rs

@@ -1,8 +1,11 @@
+use crate::common::get_indexed_attestation;
+use crate::per_block_processing::errors::{AttestationInvalid, BlockOperationError};
+use std::collections::{hash_map::Entry, HashMap};
 use std::marker::PhantomData;
 use tree_hash::TreeHash;
 use types::{
-    AbstractExecPayload, BeaconState, BeaconStateError, ChainSpec, EthSpec, Hash256,
-    SignedBeaconBlock, Slot,
+    AbstractExecPayload, Attestation, AttestationData, BeaconState, BeaconStateError, BitList,
+    ChainSpec, Epoch, EthSpec, Hash256, IndexedAttestation, SignedBeaconBlock, Slot,
 };
 
 #[derive(Debug)]
@@ -13,6 +16,9 @@ pub struct ConsensusContext<T: EthSpec> {
     proposer_index: Option<u64>,
     /// Block root of the block at `slot`.
     current_block_root: Option<Hash256>,
+    /// Cache of indexed attestations constructed during block processing.
+    indexed_attestations:
+        HashMap<(AttestationData, BitList<T::MaxValidatorsPerCommittee>), IndexedAttestation<T>>,
     _phantom: PhantomData<T>,
 }
 
@@ -20,6 +26,7 @@ pub struct ConsensusContext<T: EthSpec> {
 pub enum ContextError {
     BeaconState(BeaconStateError),
     SlotMismatch { slot: Slot, expected: Slot },
+    EpochMismatch { epoch: Epoch, expected: Epoch },
 }
 
 impl From<BeaconStateError> for ContextError {
@@ -34,6 +41,7 @@ impl<T: EthSpec> ConsensusContext<T> {
             slot,
             proposer_index: None,
             current_block_root: None,
+            indexed_attestations: HashMap::new(),
             _phantom: PhantomData,
         }
     }
@@ -43,13 +51,39 @@ impl<T: EthSpec> ConsensusContext<T> {
         self
     }
 
+    /// Strict method for fetching the proposer index.
+    ///
+    /// Gets the proposer index for `self.slot` while ensuring that it matches `state.slot()`. This
+    /// method should be used in block processing and almost everywhere the proposer index is
+    /// required. If the slot check is too restrictive, see `get_proposer_index_from_epoch_state`.
     pub fn get_proposer_index(
         &mut self,
         state: &BeaconState<T>,
         spec: &ChainSpec,
     ) -> Result<u64, ContextError> {
         self.check_slot(state.slot())?;
+        self.get_proposer_index_no_checks(state, spec)
+    }
 
+    /// More liberal method for fetching the proposer index.
+    ///
+    /// Fetches the proposer index for `self.slot` but does not require the state to be from an
+    /// exactly matching slot (merely a matching epoch). This is useful in batch verification where
+    /// we want to extract the proposer index from a single state for every slot in the epoch.
+    pub fn get_proposer_index_from_epoch_state(
+        &mut self,
+        state: &BeaconState<T>,
+        spec: &ChainSpec,
+    ) -> Result<u64, ContextError> {
+        self.check_epoch(state.current_epoch())?;
+        self.get_proposer_index_no_checks(state, spec)
+    }
+
+    fn get_proposer_index_no_checks(
+        &mut self,
+        state: &BeaconState<T>,
+        spec: &ChainSpec,
+    ) -> Result<u64, ContextError> {
         if let Some(proposer_index) = self.proposer_index {
             return Ok(proposer_index);
         }
@@ -89,4 +123,39 @@ 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 get_indexed_attestation(
+        &mut self,
+        state: &BeaconState<T>,
+        attestation: &Attestation<T>,
+    ) -> Result<&IndexedAttestation<T>, BlockOperationError<AttestationInvalid>> {
+        let key = (
+            attestation.data.clone(),
+            attestation.aggregation_bits.clone(),
+        );
+
+        match self.indexed_attestations.entry(key) {
+            Entry::Occupied(occupied) => Ok(occupied.into_mut()),
+            Entry::Vacant(vacant) => {
+                let committee =
+                    state.get_beacon_committee(attestation.data.slot, attestation.data.index)?;
+                let indexed_attestation =
+                    get_indexed_attestation(committee.committee, attestation)?;
+                Ok(vacant.insert(indexed_attestation))
+            }
+        }
+    }
+
+    pub fn num_cached_indexed_attestations(&self) -> usize {
+        self.indexed_attestations.len()
+    }
 }