perf(beacon): run epoch steps 1-3 over one registry scan

crates/blockchain/state_transition/src/beacon/helpers/altair.rs

@@ -31,7 +31,8 @@
 //! without changing the shape of any SSZ container. That is the only function
 //! below that takes a [`Config`]; the rest need nothing a network could vary.
 
-use super::finality::{get_eligible_validator_indices, is_in_inactivity_leak};
+use super::participation::{EpochSummary, RewardContext};
+use super::predicates::is_active_validator;
 use crate::beacon::bls;
 use crate::beacon::config::Config;
 use crate::beacon::constants;
@@ -45,7 +46,7 @@ use crate::beacon::primitives::{Epoch, Gwei, ParticipationFlags, ValidatorIndex}
 
 use super::accessors::{
     get_active_validator_indices, get_block_root, get_block_root_at_slot, get_current_epoch,
-    get_previous_epoch, get_seed, get_total_active_balance, get_total_balance,
+    get_previous_epoch, get_seed, get_total_active_balance,
 };
 use super::math::integer_squareroot;
 use super::shuffling::compute_shuffled_index;
@@ -212,10 +213,8 @@ pub fn get_base_reward(state: &BeaconState, index: ValidatorIndex) -> Result<Gwe
 /// and `previous_epoch_participation`, mirroring the two-epoch window phase0
 /// keeps for `PendingAttestation`s).
 ///
-/// Ascending and duplicate-free: it is built by filtering
-/// [`get_active_validator_indices`], which already returns indices in that
-/// order, so callers may binary-search it the way
-/// [`get_flag_index_deltas`] does.
+/// Ascending and duplicate-free: it is built by walking the registry in
+/// order, so callers may binary-search it.
 pub fn get_unslashed_participating_indices(
     state: &BeaconState,
     flag_index: usize,
@@ -234,18 +233,22 @@ pub fn get_unslashed_participating_indices(
         previous_epoch_participation
     };
 
+    // One pass over the registry instead of an index list followed by a
+    // descent per member.
     let mut participating_indices = Vec::new();
-    for index in get_active_validator_indices(state, epoch) {
-        let flags =
-            epoch_participation
-                .get(index as usize)
-                .copied()
-                .ok_or(Error::IndexOutOfBounds {
-                    index: index as usize,
-                    len: epoch_participation.len(),
-                })?;
-        if has_flag(flags, flag_index) && !state.validator(index)?.slashed {
-            participating_indices.push(index);
+    for (index, validator) in state.validators().iter().enumerate() {
+        if !is_active_validator(validator, epoch) {
+            continue;
+        }
+        let flags = epoch_participation
+            .get(index)
+            .copied()
+            .ok_or(Error::IndexOutOfBounds {
+                index,
+                len: epoch_participation.len(),
+            })?;
+        if has_flag(flags, flag_index) && !validator.slashed {
+            participating_indices.push(index as ValidatorIndex);
         }
     }
     Ok(participating_indices)
@@ -320,68 +323,18 @@ pub fn get_flag_index_deltas(
     state: &BeaconState,
     flag_index: usize,
 ) -> Result<(Vec<Gwei>, Vec<Gwei>)> {
-    let validator_count = state.validators().len();
-    let mut rewards = vec![0; validator_count];
-    let mut penalties = vec![0; validator_count];
-
-    let previous_epoch = get_previous_epoch(state);
-    let unslashed_participating_indices =
-        get_unslashed_participating_indices(state, flag_index, previous_epoch)?;
-    let weight = constants::PARTICIPATION_FLAG_WEIGHTS[flag_index];
-    let unslashed_participating_balance =
-        get_total_balance(state, &unslashed_participating_indices)?;
-    let unslashed_participating_increments =
-        unslashed_participating_balance / preset::EFFECTIVE_BALANCE_INCREMENT;
-    let active_increments = get_total_active_balance(state)? / preset::EFFECTIVE_BALANCE_INCREMENT;
-
-    // Hoisted out of the loop below, where the specification writes
-    // `get_base_reward(state, index)` per eligible validator. That helper is
-    // `increments * get_base_reward_per_increment(state)`, and the second
-    // factor is `get_total_active_balance`, an unconditional `O(registry
-    // size)` scan with no cache of its own. That is the same quantity
-    // `active_increments` above already paid for, just run through a
-    // different formula (`get_base_reward_per_increment` divides by
-    // `integer_squareroot`, `active_increments` does not), so it is not
-    // reusable as-is and has to be hoisted on its own.
-    //
-    // [`process_epoch::electra::process_epoch`] calls this (via
-    // `process_epoch::altair::process_rewards_and_penalties`) once per
-    // [`crate::beacon::constants::PARTICIPATION_FLAG_WEIGHTS`] entry, three times per
-    // epoch boundary. At mainnet's ~1M validators, the unhoisted form is
-    // three separate million-element scans per *eligible validator*, effectively
-    // unbounded, for what this function already computes once above. This is
-    // the same bug already fixed in `process_attestation`'s per-attester loop
-    // (see that function's own comment), left unfixed here because it runs
-    // once per epoch rather than once per block and so never showed up in a
-    // profile that did not cross an epoch boundary.
-    //
-    // Measured directly: `tests::measures_the_cost_of_get_flag_index_deltas`
-    // times this call at 2^15 validators. Unhoisted, that call took ~11.9s;
-    // hoisted, ~384us: roughly 31,000x at that scale, and the gap widens
-    // further at mainnet's ~2^20 validators, since the unhoisted form is
-    // O(n^2) (`1024x` slower again at that size) while this is O(n) (`32x`
-    // slower again, same as every other size-dependent cost in this crate).
-    let base_reward_per_increment = get_base_reward_per_increment(state)?;
-
-    for index in get_eligible_validator_indices(state) {
-        // `get_base_reward(state, index)` inlined against the hoisted
-        // per-increment value, in the helper's own order of operations so
-        // the result is bit-identical.
-        let increments =
-            state.validator(index)?.effective_balance / preset::EFFECTIVE_BALANCE_INCREMENT;
-        let base_reward = increments * base_reward_per_increment;
-        if unslashed_participating_indices
-            .binary_search(&index)
-            .is_ok()
-        {
-            if !is_in_inactivity_leak(state) {
-                let reward_numerator = base_reward * weight * unslashed_participating_increments;
-                rewards[index as usize] +=
-                    reward_numerator / (active_increments * constants::WEIGHT_DENOMINATOR);
-            }
-        } else if flag_index != constants::TIMELY_HEAD_FLAG_INDEX {
-            penalties[index as usize] += base_reward * weight / constants::WEIGHT_DENOMINATOR;
-        }
+    // Thin projection of the one-pass summary: the fixtures call this one flag
+    // at a time, so each call pays its own scan, which the epoch driver does
+    // not (it builds a single summary for all steps).
+    let summary = EpochSummary::build(state, false)?;
+    let context = RewardContext::new(state, summary.totals())?;
+
+    let mut rewards = vec![0; summary.len()];
+    let mut penalties = vec![0; summary.len()];
+    for (index, (flags, effective_balance)) in summary.iter().enumerate() {
+        let (reward, penalty) = context.flag_deltas(flags, effective_balance)[flag_index];
+        rewards[index] = reward;
+        penalties[index] = penalty;
     }
     Ok((rewards, penalties))
 }
@@ -410,39 +363,27 @@ pub fn get_inactivity_penalty_deltas(
     state: &BeaconState,
     config: &Config,
 ) -> Result<(Vec<Gwei>, Vec<Gwei>)> {
-    let validator_count = state.validators().len();
-    let rewards = vec![0; validator_count];
-    let mut penalties = vec![0; validator_count];
-
-    let previous_epoch = get_previous_epoch(state);
-    let matching_target_indices = get_unslashed_participating_indices(
-        state,
-        constants::TIMELY_TARGET_FLAG_INDEX,
-        previous_epoch,
-    )?;
-
+    let summary = EpochSummary::build(state, false)?;
+    let context = RewardContext::new(state, summary.totals())?;
     let (_, _, inactivity_scores) = state.altair_validator_lists()?;
 
-    for index in get_eligible_validator_indices(state) {
-        if matching_target_indices.binary_search(&index).is_err() {
-            let effective_balance = state.validator(index)?.effective_balance;
-            let inactivity_score =
+    let rewards = vec![0; summary.len()];
+    let mut penalties = vec![0; summary.len()];
+    for (index, (flags, effective_balance)) in summary.iter().enumerate() {
+        penalties[index] = context.inactivity_penalty(
+            flags,
+            effective_balance,
+            || {
                 inactivity_scores
-                    .get(index as usize)
+                    .get(index)
                     .copied()
                     .ok_or(Error::IndexOutOfBounds {
-                        index: index as usize,
+                        index,
                         len: inactivity_scores.len(),
-                    })?;
-
-            let penalty_numerator = effective_balance.checked_mul(inactivity_score).ok_or(
-                Error::ArithmeticOverflow("effective_balance * inactivity_scores[index]"),
-            )?;
-            let inactivity_penalty_quotient =
-                preset::retuned::inactivity_penalty_quotient(state.fork_name());
-            let penalty_denominator = config.inactivity_score_bias * inactivity_penalty_quotient;
-            penalties[index as usize] += penalty_numerator / penalty_denominator;
-        }
+                    })
+            },
+            config,
+        )?;
     }
 
     Ok((rewards, penalties))

crates/blockchain/state_transition/src/beacon/helpers/finality.rs

@@ -8,6 +8,7 @@
 //! `helpers` depend on `stf` where the dependency otherwise runs the other way.
 
 use crate::beacon::containers::BeaconState;
+use crate::beacon::error::{Error, Result};
 use crate::beacon::preset;
 use crate::beacon::primitives::{Epoch, ValidatorIndex};
 
@@ -20,8 +21,18 @@ use super::predicates::is_active_validator;
 /// by one every further epoch the chain fails to finalize, which is what lets
 /// [`is_in_inactivity_leak`] and the inactivity penalty scale with how long the
 /// stall has lasted rather than firing at a fixed severity.
-pub fn get_finality_delay(state: &BeaconState) -> Epoch {
-    get_previous_epoch(state) - state.finalized_checkpoint().epoch
+///
+/// A finalized checkpoint past the previous epoch makes this a `uint64`
+/// underflow, which the specification treats as an invalid state transition,
+/// so it returns [`Error::ArithmeticOverflow`] rather than wrapping. No chain
+/// reaches that state: justification only ever finalizes an epoch already
+/// behind the current one.
+pub fn get_finality_delay(state: &BeaconState) -> Result<Epoch> {
+    get_previous_epoch(state)
+        .checked_sub(state.finalized_checkpoint().epoch)
+        .ok_or(Error::ArithmeticOverflow(
+            "get_previous_epoch(state) - finalized_checkpoint.epoch",
+        ))
 }
 
 /// Whether the chain has gone long enough without finalizing that inactive
@@ -32,8 +43,10 @@ pub fn get_finality_delay(state: &BeaconState) -> Epoch {
 /// stalls, the faster an inactive validator's share of the active set shrinks,
 /// until the honest, active minority eventually clears the two-thirds
 /// threshold on its own.
-pub fn is_in_inactivity_leak(state: &BeaconState) -> bool {
-    get_finality_delay(state) > preset::MIN_EPOCHS_TO_INACTIVITY_PENALTY
+///
+/// Fails wherever [`get_finality_delay`] does.
+pub fn is_in_inactivity_leak(state: &BeaconState) -> Result<bool> {
+    Ok(get_finality_delay(state)? > preset::MIN_EPOCHS_TO_INACTIVITY_PENALTY)
 }
 
 /// Validators whose participation this epoch's rewards and penalties account

crates/blockchain/state_transition/src/beacon/helpers/mod.rs

@@ -21,6 +21,9 @@ pub mod fulu;
 pub mod math;
 pub mod misc;
 pub mod mutators;
+pub mod participation;
+#[cfg(any(test, debug_assertions))]
+pub mod participation_reference;
 pub mod predicates;
 pub mod shuffling;
 #[cfg(any(test, feature = "test-utils"))]