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

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 == current_epoch {
+        if initialized_epoch == 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;