Tree-ify new rewards APIs

beacon_node/beacon_chain/src/attestation_rewards.rs

@@ -42,7 +42,8 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             .ok_or(BeaconChainError::MissingBeaconState(state_root))?;
 
         // Calculate ideal_rewards
-        let participation_cache = ParticipationCache::new(&state, spec)?;
+        let participation_cache = ParticipationCache::new(&state, spec)
+            .map_err(|_| BeaconChainError::AttestationRewardsError)?;
 
         let previous_epoch = state.previous_epoch();
 
@@ -52,13 +53,9 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             let weight = get_flag_weight(flag_index)
                 .map_err(|_| BeaconChainError::AttestationRewardsError)?;
 
-            let unslashed_participating_indices = participation_cache
-                .get_unslashed_participating_indices(flag_index, previous_epoch)?;
-
-            let unslashed_participating_balance =
-                unslashed_participating_indices
-                    .total_balance()
-                    .map_err(|_| BeaconChainError::AttestationRewardsError)?;
+            let unslashed_participating_balance = participation_cache
+                .previous_epoch_flag_attesting_balance(flag_index)
+                .map_err(|_| BeaconChainError::AttestationRewardsError)?;
 
             let unslashed_participating_increments =
                 unslashed_participating_balance.safe_div(spec.effective_balance_increment)?;
@@ -110,14 +107,17 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                 .collect::<Result<Vec<_>, _>>()?
         };
 
-        for validator_index in &validators {
-            let eligible = state.is_eligible_validator(previous_epoch, *validator_index)?;
+        for &validator_index in &validators {
+            let validator = participation_cache
+                .get_validator(validator_index)
+                .map_err(|_| BeaconChainError::AttestationRewardsError)?;
+            let eligible = validator.is_eligible;
             let mut head_reward = 0u64;
             let mut target_reward = 0i64;
             let mut source_reward = 0i64;
 
             if eligible {
-                let effective_balance = state.get_effective_balance(*validator_index)?;
+                let effective_balance = validator.effective_balance;
 
                 let effective_balance_eth =
                     effective_balance.safe_div(spec.effective_balance_increment)?;
@@ -126,10 +126,8 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                     let (ideal_reward, penalty) = ideal_rewards_hashmap
                         .get(&(flag_index, effective_balance_eth))
                         .ok_or(BeaconChainError::AttestationRewardsError)?;
-                    let voted_correctly = participation_cache
-                        .get_unslashed_participating_indices(flag_index, previous_epoch)
-                        .map_err(|_| BeaconChainError::AttestationRewardsError)?
-                        .contains(*validator_index)
+                    let voted_correctly = validator
+                        .is_unslashed_participating_index(flag_index)
                         .map_err(|_| BeaconChainError::AttestationRewardsError)?;
                     if voted_correctly {
                         if flag_index == TIMELY_HEAD_FLAG_INDEX {
@@ -149,7 +147,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                 }
             }
             total_rewards.push(TotalAttestationRewards {
-                validator_index: *validator_index as u64,
+                validator_index: validator_index as u64,
                 head: head_reward,
                 target: target_reward,
                 source: source_reward,

beacon_node/beacon_chain/src/beacon_block_reward.rs

@@ -4,12 +4,11 @@ use operation_pool::RewardCache;
 use safe_arith::SafeArith;
 use slog::error;
 use state_processing::{
-    common::{
-        altair, get_attestation_participation_flag_indices, get_attesting_indices_from_state,
-    },
+    common::{get_attestation_participation_flag_indices, get_attesting_indices_from_state},
     per_block_processing::{
         altair::sync_committee::compute_sync_aggregate_rewards, get_slashable_indices,
     },
+    ConsensusContext,
 };
 use store::{
     consts::altair::{PARTICIPATION_FLAG_WEIGHTS, PROPOSER_WEIGHT, WEIGHT_DENOMINATOR},
@@ -124,7 +123,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             proposer_slashing_reward.safe_add_assign(
                 state
                     .get_validator(proposer_slashing.proposer_index() as usize)?
-                    .effective_balance
+                    .effective_balance()
                     .safe_div(self.spec.whistleblower_reward_quotient)?,
             )?;
         }
@@ -146,7 +145,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                 attester_slashing_reward.safe_add_assign(
                     state
                         .get_validator(attester_index as usize)?
-                        .effective_balance
+                        .effective_balance()
                         .safe_div(self.spec.whistleblower_reward_quotient)?,
                 )?;
             }
@@ -178,9 +177,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         block: BeaconBlockRef<'_, T::EthSpec, Payload>,
         state: &mut BeaconState<T::EthSpec>,
     ) -> Result<BeaconBlockSubRewardValue, BeaconChainError> {
-        let total_active_balance = state.get_total_active_balance()?;
-        let base_reward_per_increment =
-            altair::BaseRewardPerIncrement::new(total_active_balance, &self.spec)?;
+        let mut ctxt = ConsensusContext::new(block.slot());
 
         let mut total_proposer_reward = 0;
 
@@ -216,13 +213,9 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                     {
                         validator_participation.add_flag(flag_index)?;
                         proposer_reward_numerator.safe_add_assign(
-                            altair::get_base_reward(
-                                state,
-                                index,
-                                base_reward_per_increment,
-                                &self.spec,
-                            )?
-                            .safe_mul(weight)?,
+                            ctxt.get_base_reward(state, index, &self.spec)
+                                .map_err(|_| BeaconChainError::BlockRewardAttestationError)?
+                                .safe_mul(weight)?,
                         )?;
                     }
                 }

beacon_node/beacon_chain/src/sync_committee_rewards.rs

@@ -38,9 +38,26 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             })?;
 
         let mut balances = HashMap::<usize, u64>::new();
+        for &validator_index in &sync_committee_indices {
+            balances.insert(
+                validator_index,
+                *state
+                    .balances()
+                    .get(validator_index)
+                    .ok_or(BeaconChainError::SyncCommitteeRewardsSyncError)?,
+            );
+        }
+
+        let proposer_index = block.proposer_index() as usize;
+        balances.insert(
+            proposer_index,
+            *state
+                .balances()
+                .get(proposer_index)
+                .ok_or(BeaconChainError::SyncCommitteeRewardsSyncError)?,
+        );
 
         let mut total_proposer_rewards = 0;
-        let proposer_index = state.get_beacon_proposer_index(block.slot(), spec)?;
 
         // Apply rewards to participant balances. Keep track of proposer rewards
         for (validator_index, participant_bit) in sync_committee_indices
@@ -48,15 +65,15 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             .zip(sync_aggregate.sync_committee_bits.iter())
         {
             let participant_balance = balances
-                .entry(*validator_index)
-                .or_insert_with(|| state.balances()[*validator_index]);
+                .get_mut(&validator_index)
+                .ok_or(BeaconChainError::SyncCommitteeRewardsSyncError)?;
 
             if participant_bit {
                 participant_balance.safe_add_assign(participant_reward_value)?;
 
                 balances
-                    .entry(proposer_index)
-                    .or_insert_with(|| state.balances()[proposer_index])
+                    .get_mut(&proposer_index)
+                    .ok_or(BeaconChainError::SyncCommitteeRewardsSyncError)?
                     .safe_add_assign(proposer_reward_per_bit)?;
 
                 total_proposer_rewards.safe_add_assign(proposer_reward_per_bit)?;
@@ -67,18 +84,17 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
 
         Ok(balances
             .iter()
-            .filter_map(|(i, new_balance)| {
-                let reward = if *i != proposer_index {
-                    *new_balance as i64 - state.balances()[*i] as i64
-                } else if sync_committee_indices.contains(i) {
-                    *new_balance as i64
-                        - state.balances()[*i] as i64
-                        - total_proposer_rewards as i64
+            .filter_map(|(&i, &new_balance)| {
+                let initial_balance = *state.balances().get(i)? as i64;
+                let reward = if i != proposer_index {
+                    new_balance as i64 - initial_balance
+                } else if sync_committee_indices.contains(&i) {
+                    new_balance as i64 - initial_balance - total_proposer_rewards as i64
                 } else {
                     return None;
                 };
                 Some(SyncCommitteeReward {
-                    validator_index: *i as u64,
+                    validator_index: i as u64,
                     reward,
                 })
             })