Fix progressive balances mismatch. Perform checks only in debug mode. Some cleanups.

jimmygchen · jimmygchen · commit fef3e16d7fa6 · 2023-06-07T01:58:39.000+10:00
diff --git a/consensus/fork_choice/src/fork_choice.rs b/consensus/fork_choice/src/fork_choice.rs
@@ -13,20 +13,14 @@ use std::cmp::Ordering;
 use std::collections::BTreeSet;
 use std::marker::PhantomData;
 use std::time::Duration;
+use types::ProgressiveTotalBalances;
 use types::{
     consts::merge::INTERVALS_PER_SLOT, AbstractExecPayload, AttestationShufflingId,
     AttesterSlashing, BeaconBlockRef, BeaconState, BeaconStateError, ChainSpec, Checkpoint, Epoch,
     EthSpec, ExecPayload, ExecutionBlockHash, Hash256, IndexedAttestation, RelativeEpoch,
     SignedBeaconBlock, Slot,
 };
 
-#[cfg(test)]
-use state_processing::per_epoch_processing::{
-    weigh_justification_and_finalization, JustificationAndFinalizationState,
-};
-#[cfg(test)]
-use types::ProgressiveTotalBalances;
-
 #[derive(Debug)]
 pub enum Error<T> {
     InvalidAttestation(InvalidAttestation),
@@ -768,21 +762,17 @@ where
                         | BeaconBlockRef::Altair(_) => {
                             let participation_cache = ParticipationCache::new(state, spec)
                                 .map_err(Error::ParticipationCacheBuild)?;
-                            let processing_result =
-                            per_epoch_processing::altair::process_justification_and_finalization(
-                                state,
-                                &participation_cache,
-                            )?;
 
-                            #[cfg(test)]
-                            Self::check_processing_results_with_progressive_cache(
-                                state,
-                                &processing_result,
+                            check_progressive_total_balances(
+                                state.progressive_total_balances(),
                                 &participation_cache,
+                                spec,
                             );
 
-                            #[allow(clippy::let_and_return)]
-                            processing_result
+                            per_epoch_processing::altair::process_justification_and_finalization(
+                                state,
+                                &participation_cache,
+                            )?
                         }
                         BeaconBlockRef::Base(_) => {
                             let mut validator_statuses =
@@ -1536,59 +1526,61 @@ where
             queued_attestations: self.queued_attestations().to_vec(),
         }
     }
+}
 
-    #[cfg(test)]
-    fn check_processing_results_with_progressive_cache(
-        state: &BeaconState<E>,
-        processing_results: &JustificationAndFinalizationState<E>,
-        participation_cache: &ParticipationCache,
-    ) {
-        let expected_justified_checkpoint = processing_results.current_justified_checkpoint();
-        let expected_finalized_checkpoint = processing_results.finalized_checkpoint();
-        let expected_previous_target_balance = participation_cache
+/// Check the `ProgressiveTotalBalances` values against the computed `ParticipationCache`.
+///
+/// ## Panics  
+///
+/// In debug mode (`debug_assertions` set to `true`), this function panics if the cached progressive
+/// balances doesn't match computed values from `ParticipationCache`, or if it fails to perform the
+/// check.
+fn check_progressive_total_balances(
+    progressive_total_balances: &ProgressiveTotalBalances,
+    participation_cache: &ParticipationCache,
+    spec: &ChainSpec,
+) {
+    let do_check = || {
+        // FIXME[JC]: Converts value to 0 if it is the same as `EFFECTIVE_BALANCE_INCREMENT`.
+        // `ParticipationCache` methods return `EFFECTIVE_BALANCE_INCREMENT` (1,000,000,000)
+        // when the balance is 0, and this breaks our calculation.
+        let handle_zero_effective_balance = |val| {
+            if val == spec.effective_balance_increment {
+                0
+            } else {
+                val
+            }
+        };
+
+        let previous_target_balance = participation_cache
             .previous_epoch_target_attesting_balance()
-            .unwrap();
-        let expected_current_target_balance = participation_cache
+            .map_err(|e| BeaconStateError::ParticipationCacheError(format!("{:?}", e)))?;
+        let cached_previous_target_balance =
+            progressive_total_balances.previous_epoch_target_attesting_balance()?;
+        debug_assert!(
+            handle_zero_effective_balance(previous_target_balance) == cached_previous_target_balance,
+            "Previous epoch target attesting balance from progressive cache does not match computed value."
+        );
+
+        let current_target_balance = participation_cache
             .current_epoch_target_attesting_balance()
-            .unwrap();
-
-        // Compute justification and finalization state from progressive balances cache.
-        let (_, total_active_balance) = state
-            .total_active_balance()
-            .ok_or(BeaconStateError::TotalActiveBalanceCacheUninitialized)
-            .unwrap();
-
-        let progressive_total_balances: &ProgressiveTotalBalances =
-            state.progressive_total_balances();
-
-        let (previous_target_balance, current_target_balance) = (
-            progressive_total_balances
-                .previous_epoch_target_attesting_balance()
-                .unwrap(),
-            progressive_total_balances
-                .current_epoch_target_attesting_balance()
-                .unwrap(),
+            .map_err(|e| BeaconStateError::ParticipationCacheError(format!("{:?}", e)))?;
+        let cached_current_target_balance =
+            progressive_total_balances.current_epoch_target_attesting_balance()?;
+        debug_assert!(
+            handle_zero_effective_balance(current_target_balance) == cached_current_target_balance,
+            "Current epoch target attesting balance from progressive cache does not match computed value."
         );
 
-        let justification_and_finalization_state = weigh_justification_and_finalization(
-            JustificationAndFinalizationState::new(state),
-            total_active_balance,
-            previous_target_balance,
-            current_target_balance,
-        )
-        .unwrap();
+        Ok(())
+    };
 
-        assert_eq!(previous_target_balance, expected_previous_target_balance);
-        assert_eq!(current_target_balance, expected_current_target_balance);
-        assert_eq!(
-            justification_and_finalization_state.current_justified_checkpoint(),
-            expected_justified_checkpoint
-        );
-        assert_eq!(
-            justification_and_finalization_state.finalized_checkpoint(),
-            expected_finalized_checkpoint
-        );
-    }
+    let result: Result<(), BeaconStateError> = do_check();
+    debug_assert!(
+        result.is_ok(),
+        "An error occurred while checking progressive balances: {:?}",
+        result.err()
+    );
 }
 
 /// Helper struct that is used to encode/decode the state of the `ForkChoice` as SSZ bytes.
diff --git a/consensus/state_processing/src/common/initialize_progressive_total_balances.rs b/consensus/state_processing/src/common/initialize_progressive_total_balances.rs
@@ -1,31 +1,61 @@
-use crate::per_epoch_processing::altair::participation_cache::Error as ParticipationCacheError;
 use crate::per_epoch_processing::altair::ParticipationCache;
-use types::{BeaconState, BeaconStateError, EthSpec};
+use std::borrow::Cow;
+use types::{BeaconState, BeaconStateError, ChainSpec, EthSpec};
 
+/// Initializes the `ProgressiveTotalBalances` cache using balance values from the
+/// `ParticipationCache`. If the optional `&ParticipationCache` is not supplied, it will be computed
+/// from the `BeaconState`.
 pub fn initialize_progressive_total_balances<E: EthSpec>(
     state: &mut BeaconState<E>,
-    participation_cache: &ParticipationCache,
+    maybe_participation_cache: Option<&ParticipationCache>,
+    spec: &ChainSpec,
 ) -> Result<(), BeaconStateError> {
-    let current_epoch = state.current_epoch();
-    let progressive_total_balances = state.progressive_total_balances_mut();
-
-    let to_beacon_state_error =
-        |e: ParticipationCacheError| BeaconStateError::ParticipationCacheError(format!("{:?}", e));
-
-    let (previous_epoch_target_attesting_balance, current_epoch_target_attesting_balance) = (
-        participation_cache
-            .previous_epoch_target_attesting_balance()
-            .map_err(to_beacon_state_error)?,
-        participation_cache
-            .current_epoch_target_attesting_balance()
-            .map_err(to_beacon_state_error)?,
-    );
+    if !is_progressive_balances_enabled(state)
+        || state.progressive_total_balances().is_initialized()
+    {
+        return Ok(());
+    }
+
+    let participation_cache = match maybe_participation_cache {
+        Some(cache) => Cow::Borrowed(cache),
+        None => Cow::Owned(ParticipationCache::new(state, spec)?),
+    };
+
+    // FIXME[JC]: Converts value to 0 if it is the same as `EFFECTIVE_BALANCE_INCREMENT`.
+    // `ParticipationCache` methods return `EFFECTIVE_BALANCE_INCREMENT` (1,000,000,000)
+    // when the balance is 0, and this breaks our calculation.
+    let handle_zero_effective_balance = |val| {
+        if val == spec.effective_balance_increment {
+            0
+        } else {
+            val
+        }
+    };
+
+    let previous_epoch_target_attesting_balance = participation_cache
+        .previous_epoch_target_attesting_balance()
+        .map(handle_zero_effective_balance)
+        .map_err(|e| BeaconStateError::ParticipationCacheError(format!("{:?}", e)))?;
 
-    progressive_total_balances.initialize(
+    let current_epoch_target_attesting_balance = participation_cache
+        .current_epoch_target_attesting_balance()
+        .map(handle_zero_effective_balance)
+        .map_err(|e| BeaconStateError::ParticipationCacheError(format!("{:?}", e)))?;
+
+    let current_epoch = state.current_epoch();
+    state.progressive_total_balances_mut().initialize(
         current_epoch,
         previous_epoch_target_attesting_balance,
         current_epoch_target_attesting_balance,
     );
 
     Ok(())
 }
+
+/// `ProgressiveTotalBalances` is only enabled from `Altair` as it requires `ParticipationCache`.
+fn is_progressive_balances_enabled<E: EthSpec>(state: &BeaconState<E>) -> bool {
+    match state {
+        BeaconState::Base(_) => false,
+        BeaconState::Altair(_) | BeaconState::Merge(_) | BeaconState::Capella(_) => true,
+    }
+}
diff --git a/consensus/state_processing/src/per_block_processing.rs b/consensus/state_processing/src/per_block_processing.rs
@@ -42,7 +42,6 @@ use crate::common::decrease_balance;
 use crate::StateProcessingStrategy;
 
 use crate::common::initialize_progressive_total_balances::initialize_progressive_total_balances;
-use crate::per_epoch_processing::altair::ParticipationCache;
 #[cfg(feature = "arbitrary-fuzz")]
 use arbitrary::Arbitrary;
 
@@ -112,14 +111,11 @@ pub fn per_block_processing<T: EthSpec, Payload: AbstractExecPayload<T>>(
         .map_err(BlockProcessingError::InconsistentBlockFork)?;
 
     // Verify that the `BeaconState` instantiation matches the fork at `state.slot()`.
-    let fork_name = state
+    state
         .fork_name(spec)
         .map_err(BlockProcessingError::InconsistentStateFork)?;
 
-    if fork_name != ForkName::Base && !state.progressive_total_balances().is_initialized() {
-        let participation_cache = ParticipationCache::new(state, spec)?;
-        initialize_progressive_total_balances::<T>(state, &participation_cache)?;
-    }
+    initialize_progressive_total_balances(state, None, spec)?;
 
     let verify_signatures = match block_signature_strategy {
         BlockSignatureStrategy::VerifyBulk => {
diff --git a/consensus/state_processing/src/per_block_processing/process_operations.rs b/consensus/state_processing/src/per_block_processing/process_operations.rs
@@ -170,7 +170,10 @@ pub mod altair {
                         let validator_effective_balance = state.get_effective_balance(index)?;
                         state
                             .progressive_total_balances_mut()
-                            .on_attestation(data.target.epoch, validator_effective_balance)?;
+                            .on_new_target_attestation(
+                                data.target.epoch,
+                                validator_effective_balance,
+                            )?;
                     }
                 }
             }
diff --git a/consensus/state_processing/src/per_epoch_processing/altair.rs b/consensus/state_processing/src/per_epoch_processing/altair.rs
@@ -32,7 +32,7 @@ pub fn process_epoch<T: EthSpec>(
     // Pre-compute participating indices and total balances.
     let participation_cache = ParticipationCache::new(state, spec)?;
     let sync_committee = state.current_sync_committee()?.clone();
-    initialize_progressive_total_balances::<T>(state, &participation_cache)?;
+    initialize_progressive_total_balances::<T>(state, Some(&participation_cache), spec)?;
 
     // Justification and finalization.
     let justification_and_finalization_state =
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
@@ -53,7 +53,7 @@ impl Balance {
 }
 
 /// Caches the participation values for one epoch (either the previous or current).
-#[derive(PartialEq, Debug)]
+#[derive(PartialEq, Debug, Clone)]
 struct SingleEpochParticipationCache {
     /// Maps an active validator index to their participation flags.
     ///
@@ -173,7 +173,7 @@ impl SingleEpochParticipationCache {
 }
 
 /// Maintains a cache to be used during `altair::process_epoch`.
-#[derive(PartialEq, Debug)]
+#[derive(PartialEq, Debug, Clone)]
 pub struct ParticipationCache {
     current_epoch: Epoch,
     /// Caches information about active validators pertaining to `self.current_epoch`.
diff --git a/consensus/state_processing/src/per_epoch_processing/capella.rs b/consensus/state_processing/src/per_epoch_processing/capella.rs
@@ -28,8 +28,7 @@ pub fn process_epoch<T: EthSpec>(
     // Pre-compute participating indices and total balances.
     let participation_cache = ParticipationCache::new(state, spec)?;
     let sync_committee = state.current_sync_committee()?.clone();
-
-    initialize_progressive_total_balances::<T>(state, &participation_cache)?;
+    initialize_progressive_total_balances(state, Some(&participation_cache), spec)?;
 
     // Justification and finalization.
     let justification_and_finalization_state =
diff --git a/consensus/state_processing/src/per_epoch_processing/effective_balance_updates.rs b/consensus/state_processing/src/per_epoch_processing/effective_balance_updates.rs
@@ -7,7 +7,7 @@ use types::{BeaconStateError, EthSpec, ProgressiveTotalBalances};
 
 pub fn process_effective_balance_updates<T: EthSpec>(
     state: &mut BeaconState<T>,
-    participation_cache: Option<&ParticipationCache>,
+    maybe_participation_cache: Option<&ParticipationCache>,
     spec: &ChainSpec,
 ) -> Result<(), EpochProcessingError> {
     let hysteresis_increment = spec
@@ -32,13 +32,15 @@ pub fn process_effective_balance_updates<T: EthSpec>(
                 spec.max_effective_balance,
             );
 
-            update_progressive_balances(
-                participation_cache,
-                progressive_balances_cache,
-                index,
-                old_effective_balance,
-                new_effective_balance,
-            )?;
+            if let Some(participation_cache) = maybe_participation_cache {
+                update_progressive_balances(
+                    participation_cache,
+                    progressive_balances_cache,
+                    index,
+                    old_effective_balance,
+                    new_effective_balance,
+                )?;
+            }
 
             validator.effective_balance = new_effective_balance;
         }
@@ -47,22 +49,20 @@ pub fn process_effective_balance_updates<T: EthSpec>(
 }
 
 fn update_progressive_balances(
-    participation_cache: Option<&ParticipationCache>,
+    participation_cache: &ParticipationCache,
     progressive_balances_cache: &mut ProgressiveTotalBalances,
     index: usize,
     old_effective_balance: u64,
     new_effective_balance: u64,
 ) -> Result<(), EpochProcessingError> {
-    if let Some(participation_cache) = participation_cache {
-        if old_effective_balance != new_effective_balance {
-            let is_current_epoch_target_attester =
-                participation_cache.is_current_epoch_timely_target_attester(index)?;
-            progressive_balances_cache.on_effective_balance_change(
-                is_current_epoch_target_attester,
-                old_effective_balance,
-                new_effective_balance,
-            )?;
-        }
+    if old_effective_balance != new_effective_balance {
+        let is_current_epoch_target_attester =
+            participation_cache.is_current_epoch_timely_target_attester(index)?;
+        progressive_balances_cache.on_effective_balance_change(
+            is_current_epoch_target_attester,
+            old_effective_balance,
+            new_effective_balance,
+        )?;
     }
     Ok(())
 }
diff --git a/consensus/state_processing/src/upgrade/altair.rs b/consensus/state_processing/src/upgrade/altair.rs
@@ -1,6 +1,5 @@
 use crate::common::initialize_progressive_total_balances::initialize_progressive_total_balances;
 use crate::common::{get_attestation_participation_flag_indices, get_attesting_indices};
-use crate::per_epoch_processing::altair::ParticipationCache;
 use std::mem;
 use std::sync::Arc;
 use types::{
@@ -113,10 +112,7 @@ pub fn upgrade_to_altair<E: EthSpec>(
     // Fill in previous epoch participation from the pre state's pending attestations.
     translate_participation(&mut post, &pre.previous_epoch_attestations, spec)?;
 
-    if !post.progressive_total_balances().is_initialized() {
-        let participation_cache = ParticipationCache::new(&post, spec)?;
-        initialize_progressive_total_balances(&mut post, &participation_cache)?;
-    }
+    initialize_progressive_total_balances(&mut post, None, spec)?;
 
     // Fill in sync committees
     // Note: A duplicate committee is assigned for the current and next committee at the fork
diff --git a/consensus/types/src/progressive_total_balances.rs b/consensus/types/src/progressive_total_balances.rs
diff --git a/testing/ef_tests/src/cases/operations.rs b/testing/ef_tests/src/cases/operations.rs

Original file line number	Diff line number	Diff line change
`@@ -170,7 +170,10 @@ pub mod altair {`
`170`	`170`	`let validator_effective_balance = state.get_effective_balance(index)?;`
`171`	`171`	`state`
`172`	`172`	`.progressive_total_balances_mut()`
`173`		`- .on_attestation(data.target.epoch, validator_effective_balance)?;`
	`173`	`+ .on_new_target_attestation(`
	`174`	`+ data.target.epoch,`
	`175`	`+ validator_effective_balance,`
	`176`	`+ )?;`
`174`	`177`	`}`
`175`	`178`	`}`
`176`	`179`	`}`