[Altair] Optimization for get_unslashed_participating_indices

beacon_node/operation_pool/src/lib.rs

@@ -153,7 +153,7 @@ impl<T: EthSpec> OperationPool<T> {
             .get_cached_active_validator_indices(RelativeEpoch::Current)
             .map_err(OpPoolError::GetAttestationsTotalBalanceError)?;
         let total_active_balance = state
-            .get_total_balance(&active_indices, spec)
+            .get_total_balance(active_indices, spec)
             .map_err(OpPoolError::GetAttestationsTotalBalanceError)?;
 
         // Split attestations for the previous & current epochs, so that we
@@ -1022,7 +1022,7 @@ mod release_tests {
         let active_indices = state
             .get_cached_active_validator_indices(RelativeEpoch::Current)
             .unwrap();
-        let total_active_balance = state.get_total_balance(&active_indices, spec).unwrap();
+        let total_active_balance = state.get_total_balance(active_indices, spec).unwrap();
 
         // Set of indices covered by previous attestations in `best_attestations`.
         let mut seen_indices = BTreeSet::new();

consensus/state_processing/src/per_epoch_processing/altair/inactivity_updates.rs

@@ -12,12 +12,13 @@ pub fn process_inactivity_updates<T: EthSpec>(
     state: &mut BeaconState<T>,
     spec: &ChainSpec,
 ) -> Result<(), EpochProcessingError> {
+    let unslashed_indices = state.get_unslashed_participating_indices(
+        TIMELY_TARGET_FLAG_INDEX,
+        state.previous_epoch(),
+        spec,
+    )?;
+
     for index in state.get_eligible_validator_indices()? {
-        let unslashed_indices = state.get_unslashed_participating_indices(
-            TIMELY_TARGET_FLAG_INDEX,
-            state.previous_epoch(),
-            spec,
-        )?;
         if unslashed_indices.contains(&index) {
             let inactivity_score = state.get_inactivity_score_mut(index)?;
             if *inactivity_score > 0 {

consensus/state_processing/src/per_epoch_processing/altair/justification_and_finalization.rs

@@ -28,8 +28,8 @@ pub fn process_justification_and_finalization<T: EthSpec>(
             .as_slice(),
         spec,
     )?;
-    let previous_target_balance = state.get_total_balance(previous_indices.as_slice(), spec)?;
-    let current_target_balance = state.get_total_balance(current_indices.as_slice(), spec)?;
+    let previous_target_balance = state.get_total_balance(&previous_indices, spec)?;
+    let current_target_balance = state.get_total_balance(&current_indices, spec)?;
     weigh_justification_and_finalization(
         state,
         total_active_balance,

consensus/state_processing/src/per_epoch_processing/altair/rewards_and_penalties.rs

@@ -60,7 +60,7 @@ pub fn get_flag_index_deltas<T: EthSpec>(
         state.get_unslashed_participating_indices(flag_index, state.previous_epoch(), spec)?;
     let increment = spec.effective_balance_increment; //Factored out from balances to avoid uint64 overflow
     let unslashed_participating_increments = state
-        .get_total_balance(unslashed_participating_indices.as_slice(), spec)?
+        .get_total_balance(&unslashed_participating_indices, spec)?
         .safe_div(increment)?;
     let active_increments = total_active_balance.safe_div(increment)?;
 

consensus/types/src/beacon_state.rs

@@ -13,6 +13,7 @@ use serde_derive::{Deserialize, Serialize};
 use ssz::{ssz_encode, Decode, DecodeError, Encode};
 use ssz_derive::{Decode, Encode};
 use ssz_types::{typenum::Unsigned, BitVector, FixedVector};
+use std::collections::HashSet;
 use std::convert::TryInto;
 use std::{fmt, mem};
 use superstruct::superstruct;
@@ -1196,12 +1197,12 @@ impl<T: EthSpec> BeaconState<T> {
     /// Implementation of `get_total_balance`, matching the spec.
     ///
     /// Returns minimum `EFFECTIVE_BALANCE_INCREMENT`, to avoid div by 0.
-    pub fn get_total_balance(
-        &self,
-        validator_indices: &[usize],
+    pub fn get_total_balance<'a, I: IntoIterator<Item = &'a usize>>(
+        &'a self,
+        validator_indices: I,
         spec: &ChainSpec,
     ) -> Result<u64, Error> {
-        let total_balance = validator_indices.iter().try_fold(0_u64, |acc, i| {
+        let total_balance = validator_indices.into_iter().try_fold(0_u64, |acc, i| {
             self.get_effective_balance(*i)
                 .and_then(|bal| Ok(acc.safe_add(bal)?))
         })?;
@@ -1581,7 +1582,7 @@ impl<T: EthSpec> BeaconState<T> {
         flag_index: u32,
         epoch: Epoch,
         spec: &ChainSpec,
-    ) -> Result<Vec<usize>, Error> {
+    ) -> Result<HashSet<usize>, Error> {
         let epoch_participation = if epoch == self.current_epoch() {
             self.current_epoch_participation()?
         } else if epoch == self.previous_epoch() {