Update attestation rewards API for Electra

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

@@ -30,7 +30,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};
@@ -231,6 +231,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,
@@ -252,16 +253,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")
@@ -317,16 +324,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")
@@ -374,6 +387,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,
@@ -559,6 +573,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());