Tree states optimization using EpochCache (#4429)

* Relocate epoch cache to BeaconState

* Optimize per block processing by pulling previous epoch & current epoch calculation up.

* Revert `get_cow` change (no performance improvement)

* Initialize `EpochCache` in epoch processing and load it from state when getting base rewards.

* Initialize `EpochCache` at start of block processing if required.

* Initialize `EpochCache` in `transition_blocks` if `exclude_cache_builds` is enabled

* Fix epoch cache initialization logic

* Remove FIXME comment.

* Cache previous & current epochs in `consensus_context.rs`.

* Move `get_base_rewards` from `ConsensusContext` to `BeaconState`.

* Update Milhouse version

Cargo.lock

@@ -5160,7 +5160,7 @@ dependencies = [
 [[package]]
 name = "milhouse"
 version = "0.1.0"
-source = "git+https://github.com/sigp/milhouse?branch=main#248bc353849c113bdf078c5a81e629285c1c0589"
+source = "git+https://github.com/sigp/milhouse?branch=main#e30719aaf1e4b46e3ff1962f7370965e68fa548b"
 dependencies = [
  "arbitrary",
  "derivative",

beacon_node/beacon_chain/src/beacon_block_reward.rs

@@ -8,7 +8,6 @@ use state_processing::{
     per_block_processing::{
         altair::sync_committee::compute_sync_aggregate_rewards, get_slashable_indices,
     },
-    ConsensusContext,
 };
 use store::{
     consts::altair::{PARTICIPATION_FLAG_WEIGHTS, PROPOSER_WEIGHT, WEIGHT_DENOMINATOR},
@@ -177,8 +176,6 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         block: BeaconBlockRef<'_, T::EthSpec, Payload>,
         state: &mut BeaconState<T::EthSpec>,
     ) -> Result<BeaconBlockSubRewardValue, BeaconChainError> {
-        let mut ctxt = ConsensusContext::new(block.slot());
-
         let mut total_proposer_reward = 0;
 
         let proposer_reward_denominator = WEIGHT_DENOMINATOR
@@ -202,8 +199,13 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             for index in attesting_indices {
                 let index = index as usize;
                 for (flag_index, &weight) in PARTICIPATION_FLAG_WEIGHTS.iter().enumerate() {
-                    let epoch_participation =
+                    let previous_epoch = state.previous_epoch();
-                        state.get_epoch_participation_mut(data.target.epoch)?;
+                    let current_epoch = state.current_epoch();
+                    let epoch_participation = state.get_epoch_participation_mut(
+                        data.target.epoch,
+                        previous_epoch,
+                        current_epoch,
+                    )?;
                     let validator_participation = epoch_participation
                         .get_mut(index)
                         .ok_or(BeaconStateError::ParticipationOutOfBounds(index))?;
@@ -213,7 +215,8 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                     {
                         validator_participation.add_flag(flag_index)?;
                         proposer_reward_numerator.safe_add_assign(
-                            ctxt.get_base_reward(state, index, &self.spec)
+                            state
+                                .get_base_reward(index)
                                 .map_err(|_| BeaconChainError::BlockRewardAttestationError)?
                                 .safe_mul(weight)?,
                         )?;

consensus/state_processing/src/consensus_context.rs

@@ -1,7 +1,6 @@
 use crate::common::get_indexed_attestation;
 use crate::per_block_processing::errors::{AttestationInvalid, BlockOperationError};
-use crate::{EpochCache, EpochCacheError};
+use crate::EpochCacheError;
-use std::borrow::Cow;
 use std::collections::{hash_map::Entry, HashMap};
 use std::marker::PhantomData;
 use tree_hash::TreeHash;
@@ -14,12 +13,14 @@ use types::{
 pub struct ConsensusContext<T: EthSpec> {
     /// Slot to act as an identifier/safeguard
     slot: Slot,
+    /// Previous epoch of the `slot` precomputed for optimization purpose.
+    pub(crate) previous_epoch: Epoch,
+    /// Current epoch of the `slot` precomputed for optimization purpose.
+    pub(crate) current_epoch: Epoch,
     /// Proposer index of the block at `slot`.
     proposer_index: Option<u64>,
     /// Block root of the block at `slot`.
     current_block_root: Option<Hash256>,
-    /// Epoch cache of values that are useful for block processing that are static over an epoch.
-    epoch_cache: Option<EpochCache>,
     /// Cache of indexed attestations constructed during block processing.
     indexed_attestations:
         HashMap<(AttestationData, BitList<T::MaxValidatorsPerCommittee>), IndexedAttestation<T>>,
@@ -48,11 +49,14 @@ impl From<EpochCacheError> for ContextError {
 
 impl<T: EthSpec> ConsensusContext<T> {
     pub fn new(slot: Slot) -> Self {
+        let current_epoch = slot.epoch(T::slots_per_epoch());
+        let previous_epoch = current_epoch.saturating_sub(1u64);
         Self {
             slot,
+            previous_epoch,
+            current_epoch,
             proposer_index: None,
             current_block_root: None,
-            epoch_cache: None,
             indexed_attestations: HashMap::new(),
             _phantom: PhantomData,
         }
@@ -145,31 +149,6 @@ impl<T: EthSpec> ConsensusContext<T> {
         }
     }
 
-    pub fn set_epoch_cache(mut self, epoch_cache: EpochCache) -> Self {
-        self.epoch_cache = Some(epoch_cache);
-        self
-    }
-
-    pub fn get_base_reward(
-        &mut self,
-        state: &BeaconState<T>,
-        validator_index: usize,
-        spec: &ChainSpec,
-    ) -> Result<u64, ContextError> {
-        self.check_slot(state.slot())?;
-
-        // Build epoch cache if not already built.
-        let epoch_cache = if let Some(ref cache) = self.epoch_cache {
-            Cow::Borrowed(cache)
-        } else {
-            let cache = EpochCache::new(state, spec)?;
-            self.epoch_cache = Some(cache.clone());
-            Cow::Owned(cache)
-        };
-
-        Ok(epoch_cache.get_base_reward(validator_index)?)
-    }
-
     pub fn get_indexed_attestation(
         &mut self,
         state: &BeaconState<T>,

consensus/state_processing/src/epoch_cache.rs

@@ -1,137 +1,55 @@
-use crate::common::{
+use crate::common::altair::BaseRewardPerIncrement;
-    altair::{self, BaseRewardPerIncrement},
+use crate::common::base::SqrtTotalActiveBalance;
-    base::{self, SqrtTotalActiveBalance},
+use crate::common::{altair, base};
-};
+use types::epoch_cache::{EpochCache, EpochCacheError, EpochCacheKey};
-use safe_arith::ArithError;
+use types::{BeaconState, ChainSpec, Epoch, EthSpec, Hash256};
-use std::sync::Arc;
-use types::{BeaconState, BeaconStateError, ChainSpec, Epoch, EthSpec, Hash256, Slot};
 
-/// Cache of values which are uniquely determined at the start of an epoch.
+pub fn initialize_epoch_cache<E: EthSpec>(
-///
+    state: &mut BeaconState<E>,
-/// The values are fixed with respect to the last block of the _prior_ epoch, which we refer
+    epoch: Epoch,
-/// to as the "decision block". This cache is very similar to the `BeaconProposerCache` in that
+    spec: &ChainSpec,
-/// beacon proposers are determined at exactly the same time as the values in this cache, so
+) -> Result<(), EpochCacheError> {
-/// the keys for the two caches are identical.
+    let epoch_cache: &EpochCache = state.epoch_cache();
-#[derive(Debug, PartialEq, Eq, Clone)]
+    let decision_block_root = state
-pub struct EpochCache {
+        .proposer_shuffling_decision_root(Hash256::zero())
-    inner: Arc<Inner>,
+        .map_err(EpochCacheError::BeaconState)?;
-}
 
-#[derive(Debug, PartialEq, Eq, Clone)]
+    if epoch_cache
-struct Inner {
+        .check_validity::<E>(epoch, decision_block_root)
-    /// Unique identifier for this cache, which can be used to check its validity before use
+        .is_ok()
-    /// with any `BeaconState`.
+    {
-    key: EpochCacheKey,
+        // `EpochCache` has already been initialized and is valid, no need to initialize.
-    /// Base reward for every validator in this epoch.
+        return Ok(());
-    base_rewards: Vec<u64>,
-}
-
-#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
-pub struct EpochCacheKey {
-    pub epoch: Epoch,
-    pub decision_block_root: Hash256,
-}
-
-#[derive(Debug, PartialEq, Clone)]
-pub enum EpochCacheError {
-    IncorrectEpoch { cache: Epoch, state: Epoch },
-    IncorrectDecisionBlock { cache: Hash256, state: Hash256 },
-    ValidatorIndexOutOfBounds { validator_index: usize },
-    InvalidSlot { slot: Slot },
-    Arith(ArithError),
-    BeaconState(BeaconStateError),
-}
-
-impl From<BeaconStateError> for EpochCacheError {
-    fn from(e: BeaconStateError) -> Self {
-        Self::BeaconState(e)
-    }
-}
-
-impl From<ArithError> for EpochCacheError {
-    fn from(e: ArithError) -> Self {
-        Self::Arith(e)
-    }
-}
-
-impl EpochCache {
-    pub fn new<E: EthSpec>(
-        state: &BeaconState<E>,
-        spec: &ChainSpec,
-    ) -> Result<Self, EpochCacheError> {
-        let epoch = state.current_epoch();
-        let decision_block_root = state
-            .proposer_shuffling_decision_root(Hash256::zero())
-            .map_err(EpochCacheError::BeaconState)?;
-
-        // The cache should never be constructed at slot 0 because it should only be used for
-        // block processing (which implies slot > 0) or epoch processing (which implies slot >= 32).
-        /* FIXME(sproul): EF tests like this
-        if decision_block_root.is_zero() {
-            return Err(EpochCacheError::InvalidSlot { slot: state.slot() });
-        }
-        */
-
-        // Compute base rewards.
-        let total_active_balance = state.get_total_active_balance()?;
-        let sqrt_total_active_balance = SqrtTotalActiveBalance::new(total_active_balance);
-        let base_reward_per_increment = BaseRewardPerIncrement::new(total_active_balance, spec)?;
-
-        let mut base_rewards = Vec::with_capacity(state.validators().len());
-
-        for validator in state.validators().iter() {
-            let effective_balance = validator.effective_balance();
-
-            let base_reward = if spec
-                .altair_fork_epoch
-                .map_or(false, |altair_epoch| epoch < altair_epoch)
-            {
-                base::get_base_reward(effective_balance, sqrt_total_active_balance, spec)?
-            } else {
-                altair::get_base_reward(effective_balance, base_reward_per_increment, spec)?
-            };
-            base_rewards.push(base_reward);
-        }
-
-        Ok(Self {
-            inner: Arc::new(Inner {
-                key: EpochCacheKey {
-                    epoch,
-                    decision_block_root,
-                },
-                base_rewards,
-            }),
-        })
     }
 
-    pub fn check_validity<E: EthSpec>(
+    // Compute base rewards.
-        &self,
+    let total_active_balance = state.get_total_active_balance_at_epoch(epoch)?;
-        state: &BeaconState<E>,
+    let sqrt_total_active_balance = SqrtTotalActiveBalance::new(total_active_balance);
-    ) -> Result<(), EpochCacheError> {
+    let base_reward_per_increment = BaseRewardPerIncrement::new(total_active_balance, spec)?;
-        if self.inner.key.epoch != state.current_epoch() {
+
-            return Err(EpochCacheError::IncorrectEpoch {
+    let mut base_rewards = Vec::with_capacity(state.validators().len());
-                cache: self.inner.key.epoch,
+
-                state: state.current_epoch(),
+    for validator in state.validators().iter() {
-            });
+        let effective_balance = validator.effective_balance();
-        }
+
-        let state_decision_root = state
+        let base_reward = if spec
-            .proposer_shuffling_decision_root(Hash256::zero())
+            .altair_fork_epoch
-            .map_err(EpochCacheError::BeaconState)?;
+            .map_or(false, |altair_epoch| epoch < altair_epoch)
-        if self.inner.key.decision_block_root != state_decision_root {
+        {
-            return Err(EpochCacheError::IncorrectDecisionBlock {
+            base::get_base_reward(effective_balance, sqrt_total_active_balance, spec)?
-                cache: self.inner.key.decision_block_root,
+        } else {
-                state: state_decision_root,
+            altair::get_base_reward(effective_balance, base_reward_per_increment, spec)?
-            });
+        };
-        }
+        base_rewards.push(base_reward);
-        Ok(())
     }
 
-    #[inline]
+    *state.epoch_cache_mut() = EpochCache::new(
-    pub fn get_base_reward(&self, validator_index: usize) -> Result<u64, EpochCacheError> {
+        EpochCacheKey {
-        self.inner
+            epoch,
-            .base_rewards
+            decision_block_root,
-            .get(validator_index)
+        },
-            .copied()
+        base_rewards,
-            .ok_or(EpochCacheError::ValidatorIndexOutOfBounds { validator_index })
+    );
-    }
+
+    Ok(())
 }

consensus/state_processing/src/lib.rs

@@ -30,7 +30,6 @@ pub mod verify_operation;
 
 pub use block_replayer::{BlockReplayError, BlockReplayer, StateProcessingStrategy};
 pub use consensus_context::{ConsensusContext, ContextError};
-pub use epoch_cache::{EpochCache, EpochCacheError, EpochCacheKey};
 pub use genesis::{
     eth2_genesis_time, initialize_beacon_state_from_eth1, is_valid_genesis_state,
     process_activations,
@@ -43,4 +42,5 @@ pub use per_epoch_processing::{
     errors::EpochProcessingError, process_epoch as per_epoch_processing,
 };
 pub use per_slot_processing::{per_slot_processing, Error as SlotProcessingError};
+pub use types::{EpochCache, EpochCacheError, EpochCacheKey};
 pub use verify_operation::{SigVerifiedOp, VerifyOperation, VerifyOperationAt};

consensus/state_processing/src/per_block_processing.rs

@@ -41,6 +41,7 @@ mod verify_proposer_slashing;
 use crate::common::decrease_balance;
 use crate::StateProcessingStrategy;
 
+use crate::epoch_cache::initialize_epoch_cache;
 #[cfg(feature = "arbitrary-fuzz")]
 use arbitrary::Arbitrary;
 
@@ -114,6 +115,9 @@ pub fn per_block_processing<T: EthSpec, Payload: AbstractExecPayload<T>>(
         .fork_name(spec)
         .map_err(BlockProcessingError::InconsistentStateFork)?;
 
+    // Build epoch cache if it hasn't already been built, or if it is no longer valid
+    initialize_epoch_cache(state, state.current_epoch(), spec)?;
+
     let verify_signatures = match block_signature_strategy {
         BlockSignatureStrategy::VerifyBulk => {
             // Verify all signatures in the block at once.

consensus/state_processing/src/per_block_processing/process_operations.rs

@@ -113,6 +113,7 @@ pub mod altair {
             })
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn process_attestation<T: EthSpec>(
         state: &mut BeaconState<T>,
         attestation: &Attestation<T>,
@@ -149,18 +150,22 @@ pub mod altair {
             let index = *index as usize;
 
             for (flag_index, &weight) in PARTICIPATION_FLAG_WEIGHTS.iter().enumerate() {
-                let epoch_participation = state.get_epoch_participation_mut(data.target.epoch)?;
+                let epoch_participation = state.get_epoch_participation_mut(
-                let validator_participation = epoch_participation
+                    data.target.epoch,
-                    .get_mut(index)
+                    ctxt.previous_epoch,
-                    .ok_or(BeaconStateError::ParticipationOutOfBounds(index))?;
+                    ctxt.current_epoch,
+                )?;
 
-                if participation_flag_indices.contains(&flag_index)
+                if participation_flag_indices.contains(&flag_index) {
-                    && !validator_participation.has_flag(flag_index)?
+                    let validator_participation = epoch_participation
-                {
+                        .get_mut(index)
-                    validator_participation.add_flag(flag_index)?;
+                        .ok_or(BeaconStateError::ParticipationOutOfBounds(index))?;
-                    proposer_reward_numerator.safe_add_assign(
+
-                        ctxt.get_base_reward(state, index, spec)?.safe_mul(weight)?,
+                    if !validator_participation.has_flag(flag_index)? {
-                    )?;
+                        validator_participation.add_flag(flag_index)?;
+                        proposer_reward_numerator
+                            .safe_add_assign(state.get_base_reward(index)?.safe_mul(weight)?)?;
+                    }
                 }
             }
         }

consensus/state_processing/src/per_epoch_processing/altair.rs

@@ -1,4 +1,5 @@
 use super::{process_registry_updates, process_slashings, EpochProcessingSummary, Error};
+use crate::epoch_cache::initialize_epoch_cache;
 use crate::per_epoch_processing::{
     effective_balance_updates::process_effective_balance_updates,
     historical_roots_update::process_historical_roots_update,
@@ -75,6 +76,7 @@ pub fn process_epoch<T: EthSpec>(
 
     // Rotate the epoch caches to suit the epoch transition.
     state.advance_caches(spec)?;
+    initialize_epoch_cache(state, state.next_epoch()?, spec)?;
 
     Ok(EpochProcessingSummary::Altair {
         participation_cache,

consensus/state_processing/src/per_epoch_processing/base.rs

@@ -1,4 +1,5 @@
 use super::{process_registry_updates, process_slashings, EpochProcessingSummary, Error};
+use crate::epoch_cache::initialize_epoch_cache;
 use crate::per_epoch_processing::{
     effective_balance_updates::process_effective_balance_updates,
     historical_roots_update::process_historical_roots_update,
@@ -69,6 +70,7 @@ pub fn process_epoch<T: EthSpec>(
 
     // Rotate the epoch caches to suit the epoch transition.
     state.advance_caches(spec)?;
+    initialize_epoch_cache(state, state.next_epoch()?, spec)?;
 
     Ok(EpochProcessingSummary::Base {
         total_balances: validator_statuses.total_balances,

consensus/state_processing/src/per_epoch_processing/capella.rs

@@ -11,6 +11,7 @@ use crate::per_epoch_processing::{
 };
 use types::{BeaconState, ChainSpec, EthSpec, RelativeEpoch};
 
+use crate::epoch_cache::initialize_epoch_cache;
 pub use historical_summaries_update::process_historical_summaries_update;
 
 mod historical_summaries_update;
@@ -71,6 +72,7 @@ pub fn process_epoch<T: EthSpec>(
 
     // Rotate the epoch caches to suit the epoch transition.
     state.advance_caches(spec)?;
+    initialize_epoch_cache(state, state.next_epoch()?, spec)?;
 
     Ok(EpochProcessingSummary::Altair {
         participation_cache,

consensus/state_processing/src/per_epoch_processing/errors.rs

@@ -1,5 +1,5 @@
 use crate::per_epoch_processing::altair::participation_cache::Error as ParticipationCacheError;
-use types::{milhouse, BeaconStateError, InconsistentFork};
+use types::{milhouse, BeaconStateError, EpochCacheError, InconsistentFork};
 
 #[derive(Debug, PartialEq)]
 pub enum EpochProcessingError {
@@ -26,6 +26,7 @@ pub enum EpochProcessingError {
     InvalidFlagIndex(usize),
     ParticipationCache(ParticipationCacheError),
     MilhouseError(milhouse::Error),
+    EpochCache(EpochCacheError),
 }
 
 impl From<InclusionError> for EpochProcessingError {
@@ -64,6 +65,12 @@ impl From<milhouse::Error> for EpochProcessingError {
     }
 }
 
+impl From<EpochCacheError> for EpochProcessingError {
+    fn from(e: EpochCacheError) -> Self {
+        EpochProcessingError::EpochCache(e)
+    }
+}
+
 #[derive(Debug, PartialEq)]
 pub enum InclusionError {
     /// The validator did not participate in an attestation in this period.

consensus/state_processing/src/upgrade/altair.rs

@@ -2,8 +2,8 @@ use crate::common::{get_attestation_participation_flag_indices, get_attesting_in
 use std::mem;
 use std::sync::Arc;
 use types::{
-    BeaconState, BeaconStateAltair, BeaconStateError as Error, ChainSpec, EthSpec, Fork,
+    BeaconState, BeaconStateAltair, BeaconStateError as Error, ChainSpec, EpochCache, EthSpec,
-    ParticipationFlags, PendingAttestation, RelativeEpoch, SyncCommittee, VList,
+    Fork, ParticipationFlags, PendingAttestation, RelativeEpoch, SyncCommittee, VList,
 };
 
 /// Translate the participation information from the epoch prior to the fork into Altair's format.
@@ -104,6 +104,7 @@ pub fn upgrade_to_altair<E: EthSpec>(
         committee_caches: mem::take(&mut pre.committee_caches),
         pubkey_cache: mem::take(&mut pre.pubkey_cache),
         exit_cache: mem::take(&mut pre.exit_cache),
+        epoch_cache: EpochCache::default(),
     });
 
     // Fill in previous epoch participation from the pre state's pending attestations.

consensus/state_processing/src/upgrade/capella.rs

@@ -1,6 +1,7 @@
 use std::mem;
 use types::{
-    BeaconState, BeaconStateCapella, BeaconStateError as Error, ChainSpec, EthSpec, Fork, VList,
+    BeaconState, BeaconStateCapella, BeaconStateError as Error, ChainSpec, EpochCache, EthSpec,
+    Fork, VList,
 };
 
 /// Transform a `Merge` state into an `Capella` state.
@@ -66,6 +67,7 @@ pub fn upgrade_to_capella<E: EthSpec>(
         committee_caches: mem::take(&mut pre.committee_caches),
         pubkey_cache: mem::take(&mut pre.pubkey_cache),
         exit_cache: mem::take(&mut pre.exit_cache),
+        epoch_cache: EpochCache::default(),
     });
 
     *pre_state = post;

consensus/state_processing/src/upgrade/merge.rs

@@ -1,6 +1,6 @@
 use std::mem;
 use types::{
-    BeaconState, BeaconStateError as Error, BeaconStateMerge, ChainSpec, EthSpec,
+    BeaconState, BeaconStateError as Error, BeaconStateMerge, ChainSpec, EpochCache, EthSpec,
     ExecutionPayloadHeaderMerge, Fork,
 };
 
@@ -63,6 +63,7 @@ pub fn upgrade_to_bellatrix<E: EthSpec>(
         committee_caches: mem::take(&mut pre.committee_caches),
         pubkey_cache: mem::take(&mut pre.pubkey_cache),
         exit_cache: mem::take(&mut pre.exit_cache),
+        epoch_cache: EpochCache::default(),
     });
 
     *pre_state = post;

consensus/types/src/beacon_state.rs

@@ -29,6 +29,7 @@ pub use self::committee_cache::{
     compute_committee_index_in_epoch, compute_committee_range_in_epoch, epoch_committee_count,
     CommitteeCache,
 };
+use crate::epoch_cache::EpochCache;
 pub use eth_spec::*;
 pub use iter::BlockRootsIter;
 pub use milhouse::{interface::Interface, List as VList, List, Vector as FixedVector};
@@ -435,6 +436,13 @@ where
     #[test_random(default)]
     #[metastruct(exclude)]
     pub exit_cache: ExitCache,
+    /// Epoch cache of values that are useful for block processing that are static over an epoch.
+    #[serde(skip_serializing, skip_deserializing)]
+    #[ssz(skip_serializing, skip_deserializing)]
+    #[tree_hash(skip_hashing)]
+    #[test_random(default)]
+    #[metastruct(exclude)]
+    pub epoch_cache: EpochCache,
 }
 
 impl<T: EthSpec> BeaconState<T> {
@@ -494,6 +502,7 @@ impl<T: EthSpec> BeaconState<T> {
             ],
             pubkey_cache: PubkeyCache::default(),
             exit_cache: ExitCache::default(),
+            epoch_cache: EpochCache::default(),
         })
     }
 
@@ -1436,17 +1445,20 @@ impl<T: EthSpec> BeaconState<T> {
     ///
     /// Returns minimum `EFFECTIVE_BALANCE_INCREMENT`, to avoid div by 0.
     pub fn get_total_active_balance(&self) -> Result<u64, Error> {
+        self.get_total_active_balance_at_epoch(self.current_epoch())
+    }
+
+    pub fn get_total_active_balance_at_epoch(&self, epoch: Epoch) -> Result<u64, Error> {
         let (initialized_epoch, balance) = self
             .total_active_balance()
             .ok_or(Error::TotalActiveBalanceCacheUninitialized)?;
 
-        let current_epoch = self.current_epoch();
+        if initialized_epoch == epoch {
-        if initialized_epoch == current_epoch {
             Ok(balance)
         } else {
             Err(Error::TotalActiveBalanceCacheInconsistent {
                 initialized_epoch,
-                current_epoch,
+                current_epoch: epoch,
             })
         }
     }
@@ -1472,15 +1484,17 @@ impl<T: EthSpec> BeaconState<T> {
     pub fn get_epoch_participation_mut(
         &mut self,
         epoch: Epoch,
+        previous_epoch: Epoch,
+        current_epoch: Epoch,
     ) -> Result<&mut VList<ParticipationFlags, T::ValidatorRegistryLimit>, Error> {
-        if epoch == self.current_epoch() {
+        if epoch == current_epoch {
             match self {
                 BeaconState::Base(_) => Err(BeaconStateError::IncorrectStateVariant),
                 BeaconState::Altair(state) => Ok(&mut state.current_epoch_participation),
                 BeaconState::Merge(state) => Ok(&mut state.current_epoch_participation),
                 BeaconState::Capella(state) => Ok(&mut state.current_epoch_participation),
             }
-        } else if epoch == self.previous_epoch() {
+        } else if epoch == previous_epoch {
             match self {
                 BeaconState::Base(_) => Err(BeaconStateError::IncorrectStateVariant),
                 BeaconState::Altair(state) => Ok(&mut state.previous_epoch_participation),
@@ -1769,6 +1783,10 @@ impl<T: EthSpec> BeaconState<T> {
         Ok(sync_committee)
     }
 
+    pub fn get_base_reward(&self, validator_index: usize) -> Result<u64, EpochCacheError> {
+        self.epoch_cache().get_base_reward(validator_index)
+    }
+
     // FIXME(sproul): missing eth1 data votes, they would need a ResetListDiff
     #[allow(clippy::integer_arithmetic)]
     pub fn rebase_on(&mut self, base: &Self, spec: &ChainSpec) -> Result<(), Error> {

consensus/types/src/beacon_state/compact_state.rs

@@ -51,6 +51,7 @@ macro_rules! full_to_compact {
             committee_caches: $s.committee_caches.clone(),
             pubkey_cache: $s.pubkey_cache.clone(),
             exit_cache: $s.exit_cache.clone(),
+            epoch_cache: $s.epoch_cache.clone(),
 
             // Variant-specific fields
             $(
@@ -111,6 +112,7 @@ macro_rules! compact_to_full {
             committee_caches: $inner.committee_caches,
             pubkey_cache: $inner.pubkey_cache,
             exit_cache: $inner.exit_cache,
+            epoch_cache: $inner.epoch_cache,
 
             // Variant-specific fields
             $(

consensus/types/src/epoch_cache.rs

@@ -0,0 +1,95 @@
+use crate::{BeaconStateError, Epoch, EthSpec, Hash256, Slot};
+use safe_arith::ArithError;
+use std::sync::Arc;
+
+/// Cache of values which are uniquely determined at the start of an epoch.
+///
+/// The values are fixed with respect to the last block of the _prior_ epoch, which we refer
+/// to as the "decision block". This cache is very similar to the `BeaconProposerCache` in that
+/// beacon proposers are determined at exactly the same time as the values in this cache, so
+/// the keys for the two caches are identical.
+#[derive(Debug, PartialEq, Eq, Clone, Default, arbitrary::Arbitrary)]
+pub struct EpochCache {
+    inner: Option<Arc<Inner>>,
+}
+
+#[derive(Debug, PartialEq, Eq, Clone, arbitrary::Arbitrary)]
+struct Inner {
+    /// Unique identifier for this cache, which can be used to check its validity before use
+    /// with any `BeaconState`.
+    key: EpochCacheKey,
+    /// Base reward for every validator in this epoch.
+    base_rewards: Vec<u64>,
+}
+
+#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy, arbitrary::Arbitrary)]
+pub struct EpochCacheKey {
+    pub epoch: Epoch,
+    pub decision_block_root: Hash256,
+}
+
+#[derive(Debug, PartialEq, Clone)]
+pub enum EpochCacheError {
+    IncorrectEpoch { cache: Epoch, state: Epoch },
+    IncorrectDecisionBlock { cache: Hash256, state: Hash256 },
+    ValidatorIndexOutOfBounds { validator_index: usize },
+    InvalidSlot { slot: Slot },
+    Arith(ArithError),
+    BeaconState(BeaconStateError),
+    CacheNotInitialized,
+}
+
+impl From<BeaconStateError> for EpochCacheError {
+    fn from(e: BeaconStateError) -> Self {
+        Self::BeaconState(e)
+    }
+}
+
+impl From<ArithError> for EpochCacheError {
+    fn from(e: ArithError) -> Self {
+        Self::Arith(e)
+    }
+}
+
+impl EpochCache {
+    pub fn new(key: EpochCacheKey, base_rewards: Vec<u64>) -> EpochCache {
+        Self {
+            inner: Some(Arc::new(Inner { key, base_rewards })),
+        }
+    }
+
+    pub fn check_validity<E: EthSpec>(
+        &self,
+        current_epoch: Epoch,
+        state_decision_root: Hash256,
+    ) -> Result<(), EpochCacheError> {
+        let cache = self
+            .inner
+            .as_ref()
+            .ok_or(EpochCacheError::CacheNotInitialized)?;
+        if cache.key.epoch != current_epoch {
+            return Err(EpochCacheError::IncorrectEpoch {
+                cache: cache.key.epoch,
+                state: current_epoch,
+            });
+        }
+        if cache.key.decision_block_root != state_decision_root {
+            return Err(EpochCacheError::IncorrectDecisionBlock {
+                cache: cache.key.decision_block_root,
+                state: state_decision_root,
+            });
+        }
+        Ok(())
+    }
+
+    #[inline]
+    pub fn get_base_reward(&self, validator_index: usize) -> Result<u64, EpochCacheError> {
+        self.inner
+            .as_ref()
+            .ok_or(EpochCacheError::CacheNotInitialized)?
+            .base_rewards
+            .get(validator_index)
+            .copied()
+            .ok_or(EpochCacheError::ValidatorIndexOutOfBounds { validator_index })
+    }
+}

consensus/types/src/lib.rs

@@ -92,6 +92,7 @@ pub mod sync_subnet_id;
 pub mod validator_registration_data;
 pub mod withdrawal;
 
+pub mod epoch_cache;
 pub mod slot_data;
 #[cfg(feature = "sqlite")]
 pub mod sqlite;
@@ -126,6 +127,7 @@ pub use crate::deposit_data::DepositData;
 pub use crate::deposit_message::DepositMessage;
 pub use crate::deposit_tree_snapshot::{DepositTreeSnapshot, FinalizedExecutionBlock};
 pub use crate::enr_fork_id::EnrForkId;
+pub use crate::epoch_cache::{EpochCache, EpochCacheError, EpochCacheKey};
 pub use crate::eth1_data::Eth1Data;
 pub use crate::eth_spec::EthSpecId;
 pub use crate::execution_block_hash::ExecutionBlockHash;

lcli/src/transition_blocks.rs

@@ -72,9 +72,10 @@ use eth2::{
     BeaconNodeHttpClient, SensitiveUrl, Timeouts,
 };
 use ssz::Encode;
+use state_processing::epoch_cache::initialize_epoch_cache;
 use state_processing::{
     block_signature_verifier::BlockSignatureVerifier, per_block_processing, per_slot_processing,
-    BlockSignatureStrategy, ConsensusContext, EpochCache, StateProcessingStrategy, VerifyBlockRoot,
+    BlockSignatureStrategy, ConsensusContext, StateProcessingStrategy, VerifyBlockRoot,
 };
 use std::borrow::Cow;
 use std::fs::File;
@@ -351,16 +352,13 @@ fn do_transition<T: EthSpec>(
     let mut ctxt = if let Some(ctxt) = saved_ctxt {
         ctxt.clone()
     } else {
-        let mut ctxt = ConsensusContext::new(pre_state.slot())
+        let ctxt = ConsensusContext::new(pre_state.slot())
             .set_current_block_root(block_root)
             .set_proposer_index(block.message().proposer_index());
 
         if config.exclude_cache_builds {
-            ctxt = ctxt.set_epoch_cache(
+            let epoch = pre_state.current_epoch();
-                EpochCache::new(&pre_state, spec)
+            initialize_epoch_cache(&mut pre_state, epoch, spec).map_err(|e| format!("{e:?}"))?;
-                    .map_err(|e| format!("unable to build epoch cache: {e:?}"))?,
-            );
-            *saved_ctxt = Some(ctxt.clone());
         }
         ctxt
     };