Add initial participation cache

paulhauner · paulhauner · commit 1bb1c8fa03f9 · 2021-07-15T10:44:46.000+10:00
diff --git a/consensus/state_processing/src/per_epoch_processing.rs b/consensus/state_processing/src/per_epoch_processing.rs
@@ -6,7 +6,10 @@ use errors::EpochProcessingError as Error;
 pub use registry_updates::process_registry_updates;
 use safe_arith::SafeArith;
 pub use slashings::process_slashings;
-use types::{BeaconState, ChainSpec, EthSpec};
+use types::{
+    consts::altair::{TIMELY_HEAD_FLAG_INDEX, TIMELY_SOURCE_FLAG_INDEX, TIMELY_TARGET_FLAG_INDEX},
+    BeaconState, BeaconStateError, ChainSpec, EthSpec,
+};
 pub use weigh_justification_and_finalization::weigh_justification_and_finalization;
 
 pub mod altair;
@@ -25,7 +28,90 @@ pub mod weigh_justification_and_finalization;
 #[derive(PartialEq, Debug)]
 pub struct EpochProcessingSummary {
     pub total_balances: TotalBalances,
-    pub statuses: Vec<ValidatorStatus>,
+    pub statuses: Vec<ValidatorParticipation>,
+}
+
+#[derive(PartialEq, Debug)]
+pub struct ValidatorParticipation {
+    /*
+     * Current Epoch
+     */
+    pub is_current_epoch_timely_head_attester: bool,
+    pub is_current_epoch_timely_source_attester: bool,
+    pub is_current_epoch_timely_target_attester: bool,
+    /*
+     * Previous Epoch
+     */
+    pub is_previous_epoch_timely_head_attester: bool,
+    pub is_previous_epoch_timely_source_attester: bool,
+    pub is_previous_epoch_timely_target_attester: bool,
+    pub previous_epoch_inclusion_delay: Option<u64>,
+}
+
+impl ValidatorParticipation {
+    pub fn altair<T: EthSpec>(state: &BeaconState<T>) -> Result<Vec<Self>, BeaconStateError> {
+        let state = state.as_altair()?;
+
+        let mut participations = Vec::with_capacity(state.validators.len());
+        for i in 0..state.validators.len() {
+            let current = state
+                .current_epoch_participation
+                .get(i)
+                .ok_or(BeaconStateError::UnknownValidator(i))?;
+            let previous = state
+                .previous_epoch_participation
+                .get(i)
+                .ok_or(BeaconStateError::UnknownValidator(i))?;
+
+            participations.push(ValidatorParticipation {
+                /*
+                 * Current Epoch
+                 */
+                is_current_epoch_timely_head_attester: current.has_flag(TIMELY_HEAD_FLAG_INDEX)?,
+                is_current_epoch_timely_source_attester: current
+                    .has_flag(TIMELY_SOURCE_FLAG_INDEX)?,
+                is_current_epoch_timely_target_attester: current
+                    .has_flag(TIMELY_TARGET_FLAG_INDEX)?,
+                /*
+                 * Previous Epoch
+                 */
+                is_previous_epoch_timely_head_attester: previous
+                    .has_flag(TIMELY_HEAD_FLAG_INDEX)?,
+                is_previous_epoch_timely_source_attester: previous
+                    .has_flag(TIMELY_SOURCE_FLAG_INDEX)?,
+                is_previous_epoch_timely_target_attester: previous
+                    .has_flag(TIMELY_TARGET_FLAG_INDEX)?,
+                // Field is not relevant for Altair states.
+                previous_epoch_inclusion_delay: None,
+            })
+        }
+
+        Ok(participations)
+    }
+
+    pub fn base(validator_statuses: &[ValidatorStatus]) -> Vec<Self> {
+        validator_statuses
+            .iter()
+            .map(|v| {
+                Self {
+                    /*
+                     * Current Epoch
+                     */
+                    // FIXME(paul): avoid making this always `false`.
+                    is_current_epoch_timely_head_attester: false,
+                    is_current_epoch_timely_source_attester: v.is_current_epoch_attester,
+                    is_current_epoch_timely_target_attester: v.is_current_epoch_target_attester,
+                    /*
+                     * Previous Epoch
+                     */
+                    is_previous_epoch_timely_head_attester: v.is_previous_epoch_head_attester,
+                    is_previous_epoch_timely_source_attester: v.is_previous_epoch_attester,
+                    is_previous_epoch_timely_target_attester: v.is_previous_epoch_target_attester,
+                    previous_epoch_inclusion_delay: v.inclusion_info.map(|i| i.delay),
+                }
+            })
+            .collect()
+    }
 }
 
 /// Performs per-epoch processing on some BeaconState.
diff --git a/consensus/state_processing/src/per_epoch_processing/altair.rs b/consensus/state_processing/src/per_epoch_processing/altair.rs
@@ -4,16 +4,19 @@ use crate::per_epoch_processing::{
     historical_roots_update::process_historical_roots_update,
     resets::{process_eth1_data_reset, process_randao_mixes_reset, process_slashings_reset},
     validator_statuses::ValidatorStatuses,
+    ValidatorParticipation,
 };
 pub use inactivity_updates::process_inactivity_updates;
 pub use justification_and_finalization::process_justification_and_finalization;
+pub use participation_cache::ParticipationCache;
 pub use participation_flag_updates::process_participation_flag_updates;
 pub use rewards_and_penalties::process_rewards_and_penalties;
 pub use sync_committee_updates::process_sync_committee_updates;
 use types::{BeaconState, ChainSpec, EthSpec, RelativeEpoch};
 
 pub mod inactivity_updates;
 pub mod justification_and_finalization;
+pub mod participation_cache;
 pub mod participation_flag_updates;
 pub mod rewards_and_penalties;
 pub mod sync_committee_updates;
@@ -27,8 +30,10 @@ pub fn process_epoch<T: EthSpec>(
     state.build_committee_cache(RelativeEpoch::Current, spec)?;
     state.build_committee_cache(RelativeEpoch::Next, spec)?;
 
+    let participation_cache = ParticipationCache::altair(state, spec)?;
+
     // Justification and finalization.
-    process_justification_and_finalization(state, spec)?;
+    process_justification_and_finalization(state, &participation_cache, spec)?;
 
     process_inactivity_updates(state, spec)?;
 
@@ -77,6 +82,6 @@ pub fn process_epoch<T: EthSpec>(
 
     Ok(EpochProcessingSummary {
         total_balances: validator_statuses.total_balances,
-        statuses: validator_statuses.statuses,
+        statuses: ValidatorParticipation::altair(state)?,
     })
 }
diff --git a/consensus/state_processing/src/per_epoch_processing/altair/justification_and_finalization.rs b/consensus/state_processing/src/per_epoch_processing/altair/justification_and_finalization.rs
@@ -1,12 +1,13 @@
 use crate::per_epoch_processing::weigh_justification_and_finalization;
-use crate::per_epoch_processing::Error;
+use crate::per_epoch_processing::{altair::ParticipationCache, Error};
 use safe_arith::SafeArith;
 use types::consts::altair::TIMELY_TARGET_FLAG_INDEX;
 use types::{BeaconState, ChainSpec, EthSpec};
 
 /// Update the justified and finalized checkpoints for matching target attestations.
 pub fn process_justification_and_finalization<T: EthSpec>(
     state: &mut BeaconState<T>,
+    participation_cache: &ParticipationCache,
     spec: &ChainSpec,
 ) -> Result<(), Error> {
     if state.current_epoch() <= T::genesis_epoch().safe_add(1)? {
@@ -15,13 +16,10 @@ pub fn process_justification_and_finalization<T: EthSpec>(
 
     let previous_epoch = state.previous_epoch();
     let current_epoch = state.current_epoch();
-    let previous_indices = state.get_unslashed_participating_indices(
-        TIMELY_TARGET_FLAG_INDEX,
-        previous_epoch,
-        spec,
-    )?;
-    let current_indices =
-        state.get_unslashed_participating_indices(TIMELY_TARGET_FLAG_INDEX, current_epoch, spec)?;
+    let previous_indices = participation_cache
+        .get_unslashed_participating_indices(TIMELY_TARGET_FLAG_INDEX, previous_epoch)?;
+    let current_indices = participation_cache
+        .get_unslashed_participating_indices(TIMELY_TARGET_FLAG_INDEX, current_epoch)?;
     let total_active_balance = state.get_total_balance(
         state
             .get_active_validator_indices(current_epoch, spec)?
diff --git a/consensus/state_processing/src/per_epoch_processing/altair/participation_cache.rs b/consensus/state_processing/src/per_epoch_processing/altair/participation_cache.rs
@@ -0,0 +1,131 @@
+use safe_arith::ArithError;
+use std::collections::HashMap;
+use types::{
+    consts::altair::NUM_FLAG_INDICES, BeaconState, BeaconStateError, ChainSpec, Epoch, EthSpec,
+    ParticipationFlags,
+};
+
+pub struct ParticipationCache {
+    current_epoch: Epoch,
+    current_epoch_map: HashMap<usize, ParticipationFlags>,
+    previous_epoch: Epoch,
+    previous_epoch_map: HashMap<usize, ParticipationFlags>,
+}
+
+impl ParticipationCache {
+    pub fn altair<T: EthSpec>(
+        state: &BeaconState<T>,
+        spec: &ChainSpec,
+    ) -> Result<Self, BeaconStateError> {
+        let current_epoch = state.current_epoch();
+        let previous_epoch = state.previous_epoch();
+
+        Ok(Self {
+            current_epoch,
+            current_epoch_map: get_epoch_participation(state, current_epoch, spec)?,
+            previous_epoch,
+            previous_epoch_map: get_epoch_participation(state, previous_epoch, spec)?,
+        })
+    }
+
+    pub fn get_unslashed_participating_indices(
+        &self,
+        flag_index: usize,
+        epoch: Epoch,
+    ) -> Result<UnslashedParticipatingIndices, BeaconStateError> {
+        let map = if epoch == self.current_epoch {
+            &self.current_epoch_map
+        } else if epoch == self.previous_epoch {
+            &self.previous_epoch_map
+        } else {
+            return Err(BeaconStateError::EpochOutOfBounds);
+        };
+
+        // Note: protects the iterator.
+        if flag_index >= NUM_FLAG_INDICES {
+            return Err(ArithError::Overflow.into());
+        }
+
+        Ok(UnslashedParticipatingIndices { map, flag_index })
+    }
+}
+
+pub struct UnslashedParticipatingIndices<'a> {
+    map: &'a HashMap<usize, ParticipationFlags>,
+    flag_index: usize,
+}
+
+impl<'a> UnslashedParticipatingIndices<'a> {
+    pub fn contains(&self, val_index: usize) -> Result<bool, ArithError> {
+        self.map
+            .get(&val_index)
+            .map(|participation_flags| participation_flags.has_flag(self.flag_index))
+            .unwrap_or(Ok(false))
+    }
+}
+
+pub struct UnslashedParticipatingIndicesIter<'a> {
+    iter: std::collections::hash_map::Iter<'a, usize, ParticipationFlags>,
+    flag_index: usize,
+}
+
+impl<'a> Iterator for UnslashedParticipatingIndicesIter<'a> {
+    type Item = &'a usize;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        while let Some((val_index, participation_flags)) = self.iter.next() {
+            if participation_flags
+                .has_flag(self.flag_index)
+                .unwrap_or(false)
+            {
+                return Some(val_index);
+            }
+        }
+
+        None
+    }
+}
+
+impl<'a> IntoIterator for &'a UnslashedParticipatingIndices<'a> {
+    type Item = &'a usize;
+    type IntoIter = UnslashedParticipatingIndicesIter<'a>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        // todo: check flag is legit.
+        UnslashedParticipatingIndicesIter {
+            iter: self.map.iter(),
+            flag_index: self.flag_index,
+        }
+    }
+}
+
+fn get_epoch_participation<T: EthSpec>(
+    state: &BeaconState<T>,
+    epoch: Epoch,
+    spec: &ChainSpec,
+) -> Result<HashMap<usize, ParticipationFlags>, BeaconStateError> {
+    let epoch_participation = if epoch == state.current_epoch() {
+        state.current_epoch_participation()?
+    } else if epoch == state.previous_epoch() {
+        state.previous_epoch_participation()?
+    } else {
+        return Err(BeaconStateError::EpochOutOfBounds);
+    };
+
+    // Might be too large due to slashed boiz.
+    let active_validator_indices = state.get_active_validator_indices(epoch, spec)?;
+    let mut map = HashMap::with_capacity(active_validator_indices.len());
+
+    for val_index in active_validator_indices {
+        if !state.get_validator(val_index)?.slashed {
+            map.insert(
+                val_index,
+                *epoch_participation
+                    .get(val_index)
+                    .ok_or(BeaconStateError::ParticipationOutOfBounds(val_index))?,
+            );
+        }
+    }
+
+    Ok(map)
+}
diff --git a/consensus/state_processing/src/per_epoch_processing/base.rs b/consensus/state_processing/src/per_epoch_processing/base.rs
@@ -6,6 +6,7 @@ use crate::per_epoch_processing::{
     effective_balance_updates::process_effective_balance_updates,
     historical_roots_update::process_historical_roots_update,
     resets::{process_eth1_data_reset, process_randao_mixes_reset, process_slashings_reset},
+    ValidatorParticipation,
 };
 pub use justification_and_finalization::process_justification_and_finalization;
 pub use participation_record_updates::process_participation_record_updates;
@@ -71,6 +72,6 @@ pub fn process_epoch<T: EthSpec>(
 
     Ok(EpochProcessingSummary {
         total_balances: validator_statuses.total_balances,
-        statuses: validator_statuses.statuses,
+        statuses: ValidatorParticipation::base(&validator_statuses.statuses),
     })
 }
