From 89fc386264e9fff71d6769d1adbc68c3c02afa92 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Sun, 10 Mar 2019 13:38:57 +1100
Subject: [PATCH] Add extra checks for epoch benches finalization

---
 .../benches/epoch_processing_benches.rs       | 43 ++++++++++++++++---
 .../src/beacon_state_bencher.rs               |  9 ++--
 .../src/per_epoch_processing/attester_sets.rs |  2 +-
 3 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/eth2/state_processing/benches/epoch_processing_benches.rs b/eth2/state_processing/benches/epoch_processing_benches.rs
index 149d8f28e..6f9219658 100644
--- a/eth2/state_processing/benches/epoch_processing_benches.rs
+++ b/eth2/state_processing/benches/epoch_processing_benches.rs
@@ -33,7 +33,8 @@ pub fn epoch_processing_16k_validators(c: &mut Criterion) {
 
     let (state, _keypairs) = builder.build();
 
-    // Assert that the state has the maximum possible attestations.
+    // Assert that the state has an attestations for each committee that is able to include an
+    // attestation in the state.
     let committees_per_epoch = spec.get_epoch_committee_count(validator_count);
     let committees_per_slot = committees_per_epoch / spec.slots_per_epoch;
     let previous_epoch_attestations = committees_per_epoch;
@@ -41,18 +42,26 @@ pub fn epoch_processing_16k_validators(c: &mut Criterion) {
         committees_per_slot * (spec.slots_per_epoch - spec.min_attestation_inclusion_delay);
     assert_eq!(
         state.latest_attestations.len() as u64,
-        previous_epoch_attestations + current_epoch_attestations
+        previous_epoch_attestations + current_epoch_attestations,
+        "The state should have an attestation for each committee."
     );
 
     // Assert that each attestation in the state has full participation.
     let committee_size = validator_count / committees_per_epoch as usize;
     for a in &state.latest_attestations {
-        assert_eq!(a.aggregation_bitfield.num_set_bits(), committee_size);
+        assert_eq!(
+            a.aggregation_bitfield.num_set_bits(),
+            committee_size,
+            "Each attestation in the state should have full participation"
+        );
     }
 
     // Assert that we will run the first arm of process_rewards_and_penalities
     let epochs_since_finality = state.next_epoch(&spec) - state.finalized_epoch;
-    assert!(epochs_since_finality <= 4);
+    assert_eq!(
+        epochs_since_finality, 4,
+        "Epochs since finality should be 4"
+    );
 
     bench_epoch_processing(c, &state, &spec, "16k_validators");
 }
@@ -239,8 +248,32 @@ fn bench_epoch_processing(c: &mut Criterion, state: &BeaconState, spec: &ChainSp
         .sample_size(10),
     );
 
-    let state_clone = state.clone();
+    let mut state_clone = state.clone();
     let spec_clone = spec.clone();
+    let previous_epoch = state.previous_epoch(&spec);
+    let attesters = calculate_attester_sets(&state, &spec).unwrap();
+    let active_validator_indices = calculate_active_validator_indices(&state, &spec);
+    let current_total_balance = state.get_total_balance(&active_validator_indices[..], spec);
+    let previous_total_balance = state.get_total_balance(
+        &get_active_validator_indices(&state.validator_registry, previous_epoch)[..],
+        &spec,
+    );
+    assert_eq!(
+        state_clone.finalized_epoch, state_clone.validator_registry_update_epoch,
+        "The last registry update should be at the last finalized epoch."
+    );
+    process_justification(
+        &mut state_clone,
+        current_total_balance,
+        previous_total_balance,
+        attesters.previous_epoch_boundary.balance,
+        attesters.current_epoch_boundary.balance,
+        spec,
+    );
+    assert!(
+        state_clone.finalized_epoch > state_clone.validator_registry_update_epoch,
+        "The state should have been finalized."
+    );
     c.bench(
         &format!("epoch_process_with_caches_{}", desc),
         Benchmark::new("process_validator_registry", move |b| {
diff --git a/eth2/state_processing/benching_utils/src/beacon_state_bencher.rs b/eth2/state_processing/benching_utils/src/beacon_state_bencher.rs
index 0ee4a75e9..8ad4810e7 100644
--- a/eth2/state_processing/benching_utils/src/beacon_state_bencher.rs
+++ b/eth2/state_processing/benching_utils/src/beacon_state_bencher.rs
@@ -96,7 +96,6 @@ impl BeaconStateBencher {
         let slot = epoch.start_slot(spec.slots_per_epoch);
 
         state.slot = slot;
-        state.validator_registry_update_epoch = epoch - 1;
 
         state.previous_shuffling_epoch = epoch - 1;
         state.current_shuffling_epoch = epoch;
@@ -104,10 +103,12 @@ impl BeaconStateBencher {
         state.previous_shuffling_seed = Hash256::from_low_u64_le(0);
         state.current_shuffling_seed = Hash256::from_low_u64_le(1);
 
-        state.previous_justified_epoch = epoch - 2;
-        state.justified_epoch = epoch - 1;
+        state.previous_justified_epoch = epoch - 3;
+        state.justified_epoch = epoch - 2;
         state.justification_bitfield = u64::max_value();
-        state.finalized_epoch = epoch - 1;
+
+        state.finalized_epoch = epoch - 3;
+        state.validator_registry_update_epoch = epoch - 3;
     }
 
     /// Creates a full set of attestations for the `BeaconState`. Each attestation has full
diff --git a/eth2/state_processing/src/per_epoch_processing/attester_sets.rs b/eth2/state_processing/src/per_epoch_processing/attester_sets.rs
index 1252d8057..d82774ac2 100644
--- a/eth2/state_processing/src/per_epoch_processing/attester_sets.rs
+++ b/eth2/state_processing/src/per_epoch_processing/attester_sets.rs
@@ -13,7 +13,7 @@ impl Attesters {
         for i in additional_indices {
             self.indices.insert(*i);
         }
-        self.balance.saturating_add(additional_balance);
+        self.balance = self.balance.saturating_add(additional_balance);
     }
 }