Update op_pool to use proper rewards (#707)

* Update op_pool to use proper rewards

* Fix missing use import for tests

* Address Michael's comments

* Revert to private ValidatorStatuses

* Rename variable for clearer code

* Fix update_cover function

* Remove expect

* Add WIP test for rewards

* Use aggregation_bits instead of earliest_attestation_validators

* Use earliest attestation in test and correct typo

* Fix op_pool test thanks to @michaelsproul 's help

* Change test name

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -10,6 +10,7 @@ use crate::persisted_beacon_chain::{PersistedBeaconChain, BEACON_CHAIN_DB_KEY};
 use crate::timeout_rw_lock::TimeoutRwLock;
 use lmd_ghost::LmdGhost;
 use operation_pool::{OperationPool, PersistedOperationPool};
+use parking_lot::RwLock;
 use slog::{debug, error, info, trace, warn, Logger};
 use slot_clock::SlotClock;
 use ssz::Encode;
@@ -1514,7 +1515,11 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                 graffiti,
                 proposer_slashings: proposer_slashings.into(),
                 attester_slashings: attester_slashings.into(),
-                attestations: self.op_pool.get_attestations(&state, &self.spec).into(),
+                attestations: self
+                    .op_pool
+                    .get_attestations(&state, &self.spec)
+                    .map_err(BlockProductionError::OpPoolError)?
+                    .into(),
                 deposits,
                 voluntary_exits: self.op_pool.get_voluntary_exits(&state, &self.spec).into(),
             },

beacon_node/beacon_chain/src/errors.rs

@@ -1,5 +1,6 @@
 use crate::eth1_chain::Error as Eth1ChainError;
 use crate::fork_choice::Error as ForkChoiceError;
+use operation_pool::OpPoolError;
 use ssz_types::Error as SszTypesError;
 use state_processing::per_block_processing::errors::AttestationValidationError;
 use state_processing::BlockProcessingError;
@@ -67,6 +68,7 @@ pub enum BlockProductionError {
     BlockProcessingError(BlockProcessingError),
     Eth1ChainError(Eth1ChainError),
     BeaconStateError(BeaconStateError),
+    OpPoolError(OpPoolError),
     /// The `BeaconChain` was explicitly configured _without_ a connection to eth1, therefore it
     /// cannot produce blocks.
     NoEth1ChainConnection,

eth2/operation_pool/src/attestation.rs

@@ -1,35 +1,52 @@
 use crate::max_cover::MaxCover;
-use types::{Attestation, BeaconState, BitList, EthSpec};
+use state_processing::common::{get_attesting_indices, get_base_reward};
+use std::collections::HashMap;
+use types::{Attestation, BeaconState, BitList, ChainSpec, EthSpec};
 
 pub struct AttMaxCover<'a, T: EthSpec> {
     /// Underlying attestation.
     att: &'a Attestation<T>,
-    /// Bitfield of validators that are covered by this attestation.
+    /// Mapping of validator indices and their rewards.
-    fresh_validators: BitList<T::MaxValidatorsPerCommittee>,
+    fresh_validators_rewards: HashMap<u64, u64>,
 }
 
 impl<'a, T: EthSpec> AttMaxCover<'a, T> {
     pub fn new(
         att: &'a Attestation<T>,
-        fresh_validators: BitList<T::MaxValidatorsPerCommittee>,
+        state: &BeaconState<T>,
-    ) -> Self {
+        total_active_balance: u64,
-        Self {
+        spec: &ChainSpec,
+    ) -> Option<Self> {
+        let fresh_validators = earliest_attestation_validators(att, state);
+        let indices = get_attesting_indices(state, &att.data, &fresh_validators).ok()?;
+        let fresh_validators_rewards: HashMap<u64, u64> = indices
+            .iter()
+            .map(|i| *i as u64)
+            .flat_map(|validator_index| {
+                let reward =
+                    get_base_reward(state, validator_index as usize, total_active_balance, spec)
+                        .ok()?
+                        / spec.proposer_reward_quotient;
+                Some((validator_index, reward))
+            })
+            .collect();
+        Some(Self {
             att,
-            fresh_validators,
+            fresh_validators_rewards,
-        }
+        })
     }
 }
 
 impl<'a, T: EthSpec> MaxCover for AttMaxCover<'a, T> {
     type Object = Attestation<T>;
-    type Set = BitList<T::MaxValidatorsPerCommittee>;
+    type Set = HashMap<u64, u64>;
 
     fn object(&self) -> Attestation<T> {
         self.att.clone()
     }
 
-    fn covering_set(&self) -> &BitList<T::MaxValidatorsPerCommittee> {
+    fn covering_set(&self) -> &HashMap<u64, u64> {
-        &self.fresh_validators
+        &self.fresh_validators_rewards
     }
 
     /// Sneaky: we keep all the attestations together in one bucket, even though
@@ -40,15 +57,16 @@ impl<'a, T: EthSpec> MaxCover for AttMaxCover<'a, T> {
     fn update_covering_set(
         &mut self,
         best_att: &Attestation<T>,
-        covered_validators: &BitList<T::MaxValidatorsPerCommittee>,
+        covered_validators: &HashMap<u64, u64>,
     ) {
         if self.att.data.slot == best_att.data.slot && self.att.data.index == best_att.data.index {
-            self.fresh_validators.difference_inplace(covered_validators);
+            self.fresh_validators_rewards
+                .retain(|k, _| !covered_validators.contains_key(k))
         }
     }
 
     fn score(&self) -> usize {
-        self.fresh_validators.num_set_bits()
+        self.fresh_validators_rewards.values().sum::<u64>() as usize
     }
 }
 

eth2/operation_pool/src/lib.rs

@@ -5,7 +5,7 @@ mod persistence;
 
 pub use persistence::PersistedOperationPool;
 
-use attestation::{earliest_attestation_validators, AttMaxCover};
+use attestation::AttMaxCover;
 use attestation_id::AttestationId;
 use max_cover::maximum_cover;
 use parking_lot::RwLock;
@@ -21,8 +21,8 @@ use state_processing::per_block_processing::{
 use std::collections::{hash_map, HashMap, HashSet};
 use std::marker::PhantomData;
 use types::{
-    typenum::Unsigned, Attestation, AttesterSlashing, BeaconState, ChainSpec, EthSpec,
+    typenum::Unsigned, Attestation, AttesterSlashing, BeaconState, BeaconStateError, ChainSpec,
-    ProposerSlashing, Validator, VoluntaryExit,
+    EthSpec, ProposerSlashing, RelativeEpoch, Validator, VoluntaryExit,
 };
 
 #[derive(Default, Debug)]
@@ -38,6 +38,11 @@ pub struct OperationPool<T: EthSpec + Default> {
     _phantom: PhantomData<T>,
 }
 
+#[derive(Debug, PartialEq)]
+pub enum OpPoolError {
+    GetAttestationsTotalBalanceError(BeaconStateError),
+}
+
 impl<T: EthSpec> OperationPool<T> {
     /// Create a new operation pool.
     pub fn new() -> Self {
@@ -95,13 +100,19 @@ impl<T: EthSpec> OperationPool<T> {
         &self,
         state: &BeaconState<T>,
         spec: &ChainSpec,
-    ) -> Vec<Attestation<T>> {
+    ) -> Result<Vec<Attestation<T>>, OpPoolError> {
         // Attestations for the current fork, which may be from the current or previous epoch.
         let prev_epoch = state.previous_epoch();
         let current_epoch = state.current_epoch();
         let prev_domain_bytes = AttestationId::compute_domain_bytes(prev_epoch, state, spec);
         let curr_domain_bytes = AttestationId::compute_domain_bytes(current_epoch, state, spec);
         let reader = self.attestations.read();
+        let active_indices = state
+            .get_cached_active_validator_indices(RelativeEpoch::Current)
+            .map_err(OpPoolError::GetAttestationsTotalBalanceError)?;
+        let total_active_balance = state
+            .get_total_balance(&active_indices, spec)
+            .map_err(OpPoolError::GetAttestationsTotalBalanceError)?;
         let valid_attestations = reader
             .iter()
             .filter(|(key, _)| {
@@ -119,9 +130,12 @@ impl<T: EthSpec> OperationPool<T> {
                 )
                 .is_ok()
             })
-            .map(|att| AttMaxCover::new(att, earliest_attestation_validators(att, state)));
+            .flat_map(|att| AttMaxCover::new(att, state, total_active_balance, spec));
 
-        maximum_cover(valid_attestations, T::MaxAttestations::to_usize())
+        Ok(maximum_cover(
+            valid_attestations,
+            T::MaxAttestations::to_usize(),
+        ))
     }
 
     /// Remove attestations which are too old to be included in a block.
@@ -361,7 +375,10 @@ impl<T: EthSpec + Default> PartialEq for OperationPool<T> {
 // TODO: more tests
 #[cfg(all(test, not(debug_assertions)))]
 mod release_tests {
+    use super::attestation::earliest_attestation_validators;
     use super::*;
+    use state_processing::common::{get_attesting_indices, get_base_reward};
+    use std::collections::BTreeSet;
     use types::test_utils::*;
     use types::*;
 
@@ -522,12 +539,20 @@ mod release_tests {
 
         // Before the min attestation inclusion delay, get_attestations shouldn't return anything.
         state.slot -= 1;
-        assert_eq!(op_pool.get_attestations(state, spec).len(), 0);
+        assert_eq!(
+            op_pool
+                .get_attestations(state, spec)
+                .expect("should have attestations")
+                .len(),
+            0
+        );
 
         // Then once the delay has elapsed, we should get a single aggregated attestation.
         state.slot += spec.min_attestation_inclusion_delay;
 
-        let block_attestations = op_pool.get_attestations(state, spec);
+        let block_attestations = op_pool
+            .get_attestations(state, spec)
+            .expect("Should have block attestations");
         assert_eq!(block_attestations.len(), committees.len());
 
         let agg_att = &block_attestations[0];
@@ -684,7 +709,9 @@ mod release_tests {
         assert!(op_pool.num_attestations() > max_attestations);
 
         state.slot += spec.min_attestation_inclusion_delay;
-        let best_attestations = op_pool.get_attestations(state, spec);
+        let best_attestations = op_pool
+            .get_attestations(state, spec)
+            .expect("should have best attestations");
         assert_eq!(best_attestations.len(), max_attestations);
 
         // All the best attestations should be signed by at least `big_step_size` (4) validators.
@@ -692,4 +719,104 @@ mod release_tests {
             assert!(att.aggregation_bits.num_set_bits() >= big_step_size);
         }
     }
+
+    #[test]
+    fn attestation_rewards() {
+        let small_step_size = 2;
+        let big_step_size = 4;
+
+        let (ref mut state, ref keypairs, ref spec) =
+            attestation_test_state::<MainnetEthSpec>(big_step_size);
+
+        let op_pool = OperationPool::new();
+
+        let slot = state.slot - 1;
+        let committees = state
+            .get_beacon_committees_at_slot(slot)
+            .unwrap()
+            .into_iter()
+            .map(BeaconCommittee::into_owned)
+            .collect::<Vec<_>>();
+
+        let max_attestations = <MainnetEthSpec as EthSpec>::MaxAttestations::to_usize();
+        let target_committee_size = spec.target_committee_size as usize;
+
+        // Each validator will have a multiple of 1_000_000_000 wei.
+        // Safe from overflow unless there are about 18B validators (2^64 / 1_000_000_000).
+        for i in 0..state.validators.len() {
+            state.validators[i].effective_balance = 1_000_000_000 * i as u64;
+        }
+
+        let insert_attestations = |bc: &OwnedBeaconCommittee, step_size| {
+            for i in (0..target_committee_size).step_by(step_size) {
+                let att = signed_attestation(
+                    &bc.committee,
+                    bc.index,
+                    keypairs,
+                    i..i + step_size,
+                    slot,
+                    state,
+                    spec,
+                    if i == 0 { None } else { Some(0) },
+                );
+                op_pool.insert_attestation(att, state, spec).unwrap();
+            }
+        };
+
+        for committee in &committees {
+            assert_eq!(committee.committee.len(), target_committee_size);
+            // Attestations signed by only 2-3 validators
+            insert_attestations(committee, small_step_size);
+            // Attestations signed by 4+ validators
+            insert_attestations(committee, big_step_size);
+        }
+
+        let num_small = target_committee_size / small_step_size;
+        let num_big = target_committee_size / big_step_size;
+
+        assert_eq!(op_pool.attestations.read().len(), committees.len());
+        assert_eq!(
+            op_pool.num_attestations(),
+            (num_small + num_big) * committees.len()
+        );
+        assert!(op_pool.num_attestations() > max_attestations);
+
+        state.slot += spec.min_attestation_inclusion_delay;
+        let best_attestations = op_pool
+            .get_attestations(state, spec)
+            .expect("should have valid best attestations");
+        assert_eq!(best_attestations.len(), max_attestations);
+
+        let active_indices = state
+            .get_cached_active_validator_indices(RelativeEpoch::Current)
+            .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();
+        // Used for asserting that rewards are in decreasing order.
+        let mut prev_reward = u64::max_value();
+
+        for att in &best_attestations {
+            let fresh_validators_bitlist = earliest_attestation_validators(att, state);
+            let att_indices =
+                get_attesting_indices(state, &att.data, &fresh_validators_bitlist).unwrap();
+            let fresh_indices = &att_indices - &seen_indices;
+
+            let rewards = fresh_indices
+                .iter()
+                .map(|validator_index| {
+                    get_base_reward(state, *validator_index as usize, total_active_balance, spec)
+                        .unwrap()
+                        / spec.proposer_reward_quotient
+                })
+                .sum();
+
+            // Check that rewards are in decreasing order
+            assert!(prev_reward >= rewards);
+
+            prev_reward = rewards;
+            seen_indices.extend(fresh_indices);
+        }
+    }
 }

eth2/state_processing/src/common/get_base_reward.rs

@@ -0,0 +1,23 @@
+use integer_sqrt::IntegerSquareRoot;
+use types::*;
+
+/// Returns the base reward for some validator.
+///
+/// Spec v0.9.1
+pub fn get_base_reward<T: EthSpec>(
+    state: &BeaconState<T>,
+    index: usize,
+    // Should be == get_total_active_balance(state, spec)
+    total_active_balance: u64,
+    spec: &ChainSpec,
+) -> Result<u64, BeaconStateError> {
+    if total_active_balance == 0 {
+        Ok(0)
+    } else {
+        Ok(
+            state.get_effective_balance(index, spec)? * spec.base_reward_factor
+                / total_active_balance.integer_sqrt()
+                / spec.base_rewards_per_epoch,
+        )
+    }
+}

eth2/state_processing/src/common/mod.rs

@@ -1,9 +1,11 @@
 mod get_attesting_indices;
+mod get_base_reward;
 mod get_indexed_attestation;
 mod initiate_validator_exit;
 mod slash_validator;
 
 pub use get_attesting_indices::get_attesting_indices;
+pub use get_base_reward::get_base_reward;
 pub use get_indexed_attestation::get_indexed_attestation;
 pub use initiate_validator_exit::initiate_validator_exit;
 pub use slash_validator::slash_validator;

eth2/state_processing/src/per_epoch_processing/apply_rewards.rs

@@ -1,6 +1,7 @@
+use super::super::common::get_base_reward;
 use super::validator_statuses::{TotalBalances, ValidatorStatus, ValidatorStatuses};
 use super::Error;
-use integer_sqrt::IntegerSquareRoot;
+
 use types::*;
 
 /// Use to track the changes to a validators balance.
@@ -211,24 +212,3 @@ fn get_attestation_delta<T: EthSpec>(
 
     delta
 }
-
-/// Returns the base reward for some validator.
-///
-/// Spec v0.9.1
-fn get_base_reward<T: EthSpec>(
-    state: &BeaconState<T>,
-    index: usize,
-    // Should be == get_total_active_balance(state, spec)
-    total_active_balance: u64,
-    spec: &ChainSpec,
-) -> Result<u64, BeaconStateError> {
-    if total_active_balance == 0 {
-        Ok(0)
-    } else {
-        Ok(
-            state.get_effective_balance(index, spec)? * spec.base_reward_factor
-                / total_active_balance.integer_sqrt()
-                / spec.base_rewards_per_epoch,
-        )
-    }
-}