merge with upstream

consensus/state_processing/src/consensus_context.rs

@@ -1,9 +1,13 @@
+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 std::sync::Arc;
 use tree_hash::TreeHash;
 use types::{
-    AbstractExecPayload, BeaconState, BeaconStateError, BlobsSidecar, ChainSpec, EthSpec,
-    ExecPayload, Hash256, SignedBeaconBlock, Slot,
+    AbstractExecPayload, Attestation, AttestationData, BeaconState, BeaconStateError, BitList,
+    BlobsSidecar, ChainSpec, Epoch, EthSpec, ExecPayload, Hash256, IndexedAttestation,
+    SignedBeaconBlock, Slot,
 };
 
 #[derive(Debug)]
@@ -14,6 +18,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>>,
     /// Should only be populated if the sidecar has not been validated.
     blobs_sidecar: Option<Arc<BlobsSidecar<T>>>,
     /// Whether `validate_blobs_sidecar` has successfully passed.
@@ -26,6 +33,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 {
@@ -40,6 +48,7 @@ impl<T: EthSpec> ConsensusContext<T> {
             slot,
             proposer_index: None,
             current_block_root: None,
+            indexed_attestations: HashMap::new(),
             blobs_sidecar: None,
             blobs_sidecar_validated: false,
             blobs_verified_vs_txs: false,
@@ -51,13 +60,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);
         }
@@ -98,6 +133,41 @@ 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()
+    }
+
     pub fn set_blobs_sidecar_validated(mut self, blobs_sidecar_validated: bool) -> Self {
         self.blobs_sidecar_validated = blobs_sidecar_validated;
         self