Update attestation rewards API for Electra (#6819)

Closes:

- #6818


  Use `MAX_EFFECTIVE_BALANCE_ELECTRA` (2048) for attestation reward calculations involving Electra.

Add a new `InteropGenesisBuilder` that tries to provide a more flexible way to build genesis states. Unfortunately due to lifetime jank, it is quite unergonomic at present. We may want to refactor this builder in future to make it easier to use.

Author: michaelsproul

beacon_node/beacon_chain/src/attestation_rewards.rs

@@ -175,7 +175,9 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             let base_reward_per_increment =
                 BaseRewardPerIncrement::new(total_active_balance, spec)?;
 
-            for effective_balance_eth in 1..=self.max_effective_balance_increment_steps()? {
+            for effective_balance_eth in
+                1..=self.max_effective_balance_increment_steps(previous_epoch)?
+            {
                 let effective_balance =
                     effective_balance_eth.safe_mul(spec.effective_balance_increment)?;
                 let base_reward =
@@ -321,11 +323,14 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         })
     }
 
-    fn max_effective_balance_increment_steps(&self) -> Result<u64, BeaconChainError> {
+    fn max_effective_balance_increment_steps(
+        &self,
+        rewards_epoch: Epoch,
+    ) -> Result<u64, BeaconChainError> {
         let spec = &self.spec;
-        let max_steps = spec
-            .max_effective_balance
-            .safe_div(spec.effective_balance_increment)?;
+        let fork_name = spec.fork_name_at_epoch(rewards_epoch);
+        let max_effective_balance = spec.max_effective_balance_for_fork(fork_name);
+        let max_steps = max_effective_balance.safe_div(spec.effective_balance_increment)?;
         Ok(max_steps)
     }
 
@@ -386,7 +391,9 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
 
         let mut ideal_attestation_rewards_list = Vec::new();
         let sqrt_total_active_balance = SqrtTotalActiveBalance::new(total_balances.current_epoch());
-        for effective_balance_step in 1..=self.max_effective_balance_increment_steps()? {
+        for effective_balance_step in
+            1..=self.max_effective_balance_increment_steps(previous_epoch)?
+        {
             let effective_balance =
                 effective_balance_step.safe_mul(spec.effective_balance_increment)?;
             let base_reward =

beacon_node/beacon_chain/src/test_utils.rs

@@ -31,7 +31,7 @@ use execution_layer::{
     ExecutionLayer,
 };
 use futures::channel::mpsc::Receiver;
-pub use genesis::{interop_genesis_state_with_eth1, DEFAULT_ETH1_BLOCK_HASH};
+pub use genesis::{InteropGenesisBuilder, DEFAULT_ETH1_BLOCK_HASH};
 use int_to_bytes::int_to_bytes32;
 use kzg::trusted_setup::get_trusted_setup;
 use kzg::{Kzg, TrustedSetup};
@@ -232,6 +232,7 @@ pub struct Builder<T: BeaconChainTypes> {
     mock_execution_layer: Option<MockExecutionLayer<T::EthSpec>>,
     testing_slot_clock: Option<TestingSlotClock>,
     validator_monitor_config: Option<ValidatorMonitorConfig>,
+    genesis_state_builder: Option<InteropGenesisBuilder<T::EthSpec>>,
     import_all_data_columns: bool,
     runtime: TestRuntime,
     log: Logger,
@@ -253,16 +254,22 @@ impl<E: EthSpec> Builder<EphemeralHarnessType<E>> {
             )
             .unwrap(),
         );
+        let genesis_state_builder = self.genesis_state_builder.take().unwrap_or_else(|| {
+            // Set alternating withdrawal credentials if no builder is specified.
+            InteropGenesisBuilder::default().set_alternating_eth1_withdrawal_credentials()
+        });
+
         let mutator = move |builder: BeaconChainBuilder<_>| {
             let header = generate_genesis_header::<E>(builder.get_spec(), false);
-            let genesis_state = interop_genesis_state_with_eth1::<E>(
-                &validator_keypairs,
-                HARNESS_GENESIS_TIME,
-                Hash256::from_slice(DEFAULT_ETH1_BLOCK_HASH),
-                header,
-                builder.get_spec(),
-            )
-            .expect("should generate interop state");
+            let genesis_state = genesis_state_builder
+                .set_opt_execution_payload_header(header)
+                .build_genesis_state(
+                    &validator_keypairs,
+                    HARNESS_GENESIS_TIME,
+                    Hash256::from_slice(DEFAULT_ETH1_BLOCK_HASH),
+                    builder.get_spec(),
+                )
+                .expect("should generate interop state");
             builder
                 .genesis_state(genesis_state)
                 .expect("should build state using recent genesis")
@@ -318,16 +325,22 @@ impl<E: EthSpec> Builder<DiskHarnessType<E>> {
             .clone()
             .expect("cannot build without validator keypairs");
 
+        let genesis_state_builder = self.genesis_state_builder.take().unwrap_or_else(|| {
+            // Set alternating withdrawal credentials if no builder is specified.
+            InteropGenesisBuilder::default().set_alternating_eth1_withdrawal_credentials()
+        });
+
         let mutator = move |builder: BeaconChainBuilder<_>| {
             let header = generate_genesis_header::<E>(builder.get_spec(), false);
-            let genesis_state = interop_genesis_state_with_eth1::<E>(
-                &validator_keypairs,
-                HARNESS_GENESIS_TIME,
-                Hash256::from_slice(DEFAULT_ETH1_BLOCK_HASH),
-                header,
-                builder.get_spec(),
-            )
-            .expect("should generate interop state");
+            let genesis_state = genesis_state_builder
+                .set_opt_execution_payload_header(header)
+                .build_genesis_state(
+                    &validator_keypairs,
+                    HARNESS_GENESIS_TIME,
+                    Hash256::from_slice(DEFAULT_ETH1_BLOCK_HASH),
+                    builder.get_spec(),
+                )
+                .expect("should generate interop state");
             builder
                 .genesis_state(genesis_state)
                 .expect("should build state using recent genesis")
@@ -375,6 +388,7 @@ where
             mock_execution_layer: None,
             testing_slot_clock: None,
             validator_monitor_config: None,
+            genesis_state_builder: None,
             import_all_data_columns: false,
             runtime,
             log,
@@ -560,6 +574,15 @@ where
         self
     }
 
+    pub fn with_genesis_state_builder(
+        mut self,
+        f: impl FnOnce(InteropGenesisBuilder<E>) -> InteropGenesisBuilder<E>,
+    ) -> Self {
+        let builder = self.genesis_state_builder.take().unwrap_or_default();
+        self.genesis_state_builder = Some(f(builder));
+        self
+    }
+
     pub fn build(self) -> BeaconChainHarness<BaseHarnessType<E, Hot, Cold>> {
         let (shutdown_tx, shutdown_receiver) = futures::channel::mpsc::channel(1);
 

beacon_node/beacon_chain/tests/rewards.rs

@@ -36,6 +36,38 @@ fn get_harness(spec: ChainSpec) -> BeaconChainHarness<EphemeralHarnessType<E>> {
         .keypairs(KEYPAIRS.to_vec())
         .fresh_ephemeral_store()
         .chain_config(chain_config)
+        .mock_execution_layer()
+        .build();
+
+    harness.advance_slot();
+
+    harness
+}
+
+fn get_electra_harness(spec: ChainSpec) -> BeaconChainHarness<EphemeralHarnessType<E>> {
+    let chain_config = ChainConfig {
+        reconstruct_historic_states: true,
+        ..Default::default()
+    };
+
+    let spec = Arc::new(spec);
+
+    let harness = BeaconChainHarness::builder(E::default())
+        .spec(spec.clone())
+        .keypairs(KEYPAIRS.to_vec())
+        .with_genesis_state_builder(|builder| {
+            builder.set_initial_balance_fn(Box::new(move |i| {
+                // Use a variety of balances between min activation balance and max effective balance.
+                let balance = spec.max_effective_balance_electra
+                    / (i as u64 + 1)
+                    / spec.effective_balance_increment
+                    * spec.effective_balance_increment;
+                balance.max(spec.min_activation_balance)
+            }))
+        })
+        .fresh_ephemeral_store()
+        .chain_config(chain_config)
+        .mock_execution_layer()
         .build();
 
     harness.advance_slot();
@@ -560,6 +592,83 @@ async fn test_rewards_altair_inactivity_leak_justification_epoch() {
     assert_eq!(expected_balances, balances);
 }
 
+#[tokio::test]
+async fn test_rewards_electra() {
+    let spec = ForkName::Electra.make_genesis_spec(E::default_spec());
+    let harness = get_electra_harness(spec.clone());
+    let target_epoch = 0;
+
+    // advance until epoch N + 1 and get initial balances
+    harness
+        .extend_slots((E::slots_per_epoch() * (target_epoch + 1)) as usize)
+        .await;
+    let mut expected_balances = harness.get_current_state().balances().to_vec();
+
+    // advance until epoch N + 2 and build proposal rewards map
+    let mut proposal_rewards_map = HashMap::new();
+    let mut sync_committee_rewards_map = HashMap::new();
+    for _ in 0..E::slots_per_epoch() {
+        let state = harness.get_current_state();
+        let slot = state.slot() + Slot::new(1);
+
+        // calculate beacon block rewards / penalties
+        let ((signed_block, _maybe_blob_sidecars), mut state) =
+            harness.make_block_return_pre_state(state, slot).await;
+        let beacon_block_reward = harness
+            .chain
+            .compute_beacon_block_reward(signed_block.message(), &mut state)
+            .unwrap();
+
+        let total_proposer_reward = proposal_rewards_map
+            .entry(beacon_block_reward.proposer_index)
+            .or_insert(0);
+        *total_proposer_reward += beacon_block_reward.total as i64;
+
+        // calculate sync committee rewards / penalties
+        let reward_payload = harness
+            .chain
+            .compute_sync_committee_rewards(signed_block.message(), &mut state)
+            .unwrap();
+
+        for reward in reward_payload {
+            let total_sync_reward = sync_committee_rewards_map
+                .entry(reward.validator_index)
+                .or_insert(0);
+            *total_sync_reward += reward.reward;
+        }
+
+        harness.extend_slots(1).await;
+    }
+
+    // compute reward deltas for all validators in epoch N
+    let StandardAttestationRewards {
+        ideal_rewards,
+        total_rewards,
+    } = harness
+        .chain
+        .compute_attestation_rewards(Epoch::new(target_epoch), vec![])
+        .unwrap();
+
+    // assert ideal rewards are greater than 0
+    assert_eq!(
+        ideal_rewards.len() as u64,
+        spec.max_effective_balance_electra / spec.effective_balance_increment
+    );
+    assert!(ideal_rewards
+        .iter()
+        .all(|reward| reward.head > 0 && reward.target > 0 && reward.source > 0));
+
+    // apply attestation, proposal, and sync committee rewards and penalties to initial balances
+    apply_attestation_rewards(&mut expected_balances, total_rewards);
+    apply_other_rewards(&mut expected_balances, &proposal_rewards_map);
+    apply_other_rewards(&mut expected_balances, &sync_committee_rewards_map);
+
+    // verify expected balances against actual balances
+    let balances: Vec<u64> = harness.get_current_state().balances().to_vec();
+
+    assert_eq!(expected_balances, balances);
+}
+
 #[tokio::test]
 async fn test_rewards_base_subset_only() {
     let spec = ForkName::Base.make_genesis_spec(E::default_spec());