From 723283bd0125218c4e7852203dcb097f6dc88fd2 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Fri, 21 Jun 2019 11:44:15 +1000
Subject: [PATCH] Add attestation strategy to chain harness

---
 beacon_node/beacon_chain/src/test_utils.rs | 111 +++++++++++++--------
 1 file changed, 70 insertions(+), 41 deletions(-)

diff --git a/beacon_node/beacon_chain/src/test_utils.rs b/beacon_node/beacon_chain/src/test_utils.rs
index b6873fa571..65be6fba54 100644
--- a/beacon_node/beacon_chain/src/test_utils.rs
+++ b/beacon_node/beacon_chain/src/test_utils.rs
@@ -20,6 +20,11 @@ pub enum BuildStrategy {
     ForkCanonicalChainAt(Slot),
 }
 
+#[derive(Clone, Copy, Debug)]
+pub enum AttestationStrategy {
+    AllValidators,
+}
+
 pub struct CommonTypes<L, E>
 where
     L: LmdGhost<MemoryStore, E>,
@@ -95,7 +100,12 @@ where
         self.chain.catchup_state().expect("should catchup state");
     }
 
-    pub fn extend_chain(&self, build_strategy: BuildStrategy, blocks: usize) {
+    pub fn extend_chain(
+        &self,
+        build_strategy: BuildStrategy,
+        blocks: usize,
+        attestation_strategy: AttestationStrategy,
+    ) {
         // Get an initial state to build the block upon, based on the build strategy.
         let mut state = match build_strategy {
             BuildStrategy::OnCanonicalHead => self.chain.current_state().clone(),
@@ -134,7 +144,12 @@ where
                 .expect("should not error during block processing");
 
             if let BlockProcessingOutcome::Processed { block_root } = outcome {
-                self.add_attestations_to_op_pool(&new_state, block_root, slot);
+                self.add_attestations_to_op_pool(
+                    attestation_strategy,
+                    &new_state,
+                    block_root,
+                    slot,
+                );
             } else {
                 panic!("block should be successfully processed");
             }
@@ -198,6 +213,7 @@ where
 
     fn add_attestations_to_op_pool(
         &self,
+        attestation_strategy: AttestationStrategy,
         state: &BeaconState<E>,
         head_block_root: Hash256,
         head_block_slot: Slot,
@@ -205,6 +221,10 @@ where
         let spec = &self.spec;
         let fork = &state.fork;
 
+        let attesting_validators: Vec<usize> = match attestation_strategy {
+            AttestationStrategy::AllValidators => (0..self.keypairs.len()).collect(),
+        };
+
         state
             .get_crosslink_committees_at_slot(state.slot)
             .expect("should get committees")
@@ -213,52 +233,57 @@ where
                 let committee_size = cc.committee.len();
 
                 for (i, validator_index) in cc.committee.iter().enumerate() {
-                    let data = self
-                        .chain
-                        .produce_attestation_data_for_block(
-                            cc.shard,
-                            head_block_root,
-                            head_block_slot,
-                            state,
-                        )
-                        .expect("should produce attestation data");
+                    // Note: searching this array is worst-case `O(n)`. A hashset could be a better
+                    // alternative.
+                    if attesting_validators.contains(validator_index) {
+                        let data = self
+                            .chain
+                            .produce_attestation_data_for_block(
+                                cc.shard,
+                                head_block_root,
+                                head_block_slot,
+                                state,
+                            )
+                            .expect("should produce attestation data");
 
-                    let mut aggregation_bitfield = Bitfield::new();
-                    aggregation_bitfield.set(i, true);
-                    aggregation_bitfield.set(committee_size, false);
+                        let mut aggregation_bitfield = Bitfield::new();
+                        aggregation_bitfield.set(i, true);
+                        aggregation_bitfield.set(committee_size, false);
 
-                    let mut custody_bitfield = Bitfield::new();
-                    custody_bitfield.set(committee_size, false);
+                        let mut custody_bitfield = Bitfield::new();
+                        custody_bitfield.set(committee_size, false);
 
-                    let signature = {
-                        let message = AttestationDataAndCustodyBit {
-                            data: data.clone(),
-                            custody_bit: false,
-                        }
-                        .tree_hash_root();
+                        let signature = {
+                            let message = AttestationDataAndCustodyBit {
+                                data: data.clone(),
+                                custody_bit: false,
+                            }
+                            .tree_hash_root();
 
-                        let domain = spec.get_domain(data.target_epoch, Domain::Attestation, fork);
+                            let domain =
+                                spec.get_domain(data.target_epoch, Domain::Attestation, fork);
 
-                        let mut agg_sig = AggregateSignature::new();
-                        agg_sig.add(&Signature::new(
-                            &message,
-                            domain,
-                            self.get_sk(*validator_index),
-                        ));
+                            let mut agg_sig = AggregateSignature::new();
+                            agg_sig.add(&Signature::new(
+                                &message,
+                                domain,
+                                self.get_sk(*validator_index),
+                            ));
 
-                        agg_sig
-                    };
+                            agg_sig
+                        };
 
-                    let attestation = Attestation {
-                        aggregation_bitfield,
-                        data,
-                        custody_bitfield,
-                        signature,
-                    };
+                        let attestation = Attestation {
+                            aggregation_bitfield,
+                            data,
+                            custody_bitfield,
+                            signature,
+                        };
 
-                    self.chain
-                        .process_attestation(attestation)
-                        .expect("should process attestation");
+                        self.chain
+                            .process_attestation(attestation)
+                            .expect("should process attestation");
+                    }
                 }
             });
     }
@@ -295,7 +320,11 @@ mod test {
 
         let harness = get_harness(VALIDATOR_COUNT);
 
-        harness.extend_chain(BuildStrategy::OnCanonicalHead, num_blocks_produced as usize);
+        harness.extend_chain(
+            BuildStrategy::OnCanonicalHead,
+            num_blocks_produced as usize,
+            AttestationStrategy::AllValidators,
+        );
 
         let state = &harness.chain.head().beacon_state;