Refactored attestation creation in test utils

This commit is contained in:
Grant Wuerker 2019-07-26 12:48:17 +02:00
parent edd99fafb6
commit 78f3911522
No known key found for this signature in database
GPG Key ID: F7EA56FDDA6C464F

View File

@ -8,11 +8,7 @@ use std::sync::Arc;
use store::MemoryStore;
use store::Store;
use tree_hash::{SignedRoot, TreeHash};
use types::{
test_utils::TestingBeaconStateBuilder, AggregateSignature, Attestation,
AttestationDataAndCustodyBit, BeaconBlock, BeaconState, Bitfield, ChainSpec, Domain, EthSpec,
Hash256, Keypair, RelativeEpoch, SecretKey, Signature, Slot,
};
use types::{test_utils::TestingBeaconStateBuilder, AggregateSignature, Attestation, AttestationDataAndCustodyBit, BeaconBlock, BeaconState, Bitfield, ChainSpec, Domain, EthSpec, Hash256, Keypair, RelativeEpoch, SecretKey, Signature, Slot, CrosslinkCommittee};
pub use crate::persisted_beacon_chain::{PersistedBeaconChain, BEACON_CHAIN_DB_KEY};
@ -171,7 +167,7 @@ where
if let BlockProcessingOutcome::Processed { block_root } = outcome {
head_block_root = Some(block_root);
self.add_attestations_to_op_pool(
self.add_attestations_to_chain(
&attestation_strategy,
&new_state,
block_root,
@ -256,18 +252,16 @@ where
(block, state)
}
/// Adds attestations to the `BeaconChain` operations pool to be included in future blocks.
/// Adds attestations to the `BeaconChain` operations pool and fork choice.
///
/// The `attestation_strategy` dictates which validators should attest.
fn add_attestations_to_op_pool(
fn add_attestations_to_chain(
&self,
attestation_strategy: &AttestationStrategy,
state: &BeaconState<E>,
head_block_root: Hash256,
head_block_slot: Slot,
) {
let spec = &self.spec;
let fork = &state.fork;
let attesting_validators: Vec<usize> = match attestation_strategy {
AttestationStrategy::AllValidators => (0..self.keypairs.len()).collect(),
@ -279,16 +273,45 @@ where
.expect("should get committees")
.iter()
.for_each(|cc| {
let committee_size = cc.committee.len();
for (i, validator_index) in cc.committee.iter().enumerate() {
// Note: searching this array is worst-case `O(n)`. A hashset could be a better
// alternative.
if attesting_validators.contains(validator_index) {
let attestation = self.create_attestation(
*validator_index,
cc,
head_block_root,
head_block_slot,
state,
i
);
self.chain
.process_attestation(attestation)
.expect("should process attestation");
}
}
});
}
/// Creates an attestation for a validator with the given data.
pub fn create_attestation(
&self,
validator_index: usize,
crosslink_committee: &CrosslinkCommittee,
head_block_root: Hash256,
head_block_slot: Slot,
state: &BeaconState<E>,
bitfield_index: usize
) -> Attestation {
let committee_size = crosslink_committee.committee.len();
let spec = &self.spec;
let fork = &state.fork;
let data = self
.chain
.produce_attestation_data_for_block(
cc.shard,
crosslink_committee.shard,
head_block_root,
head_block_slot,
state,
@ -296,7 +319,7 @@ where
.expect("should produce attestation data");
let mut aggregation_bitfield = Bitfield::new();
aggregation_bitfield.set(i, true);
aggregation_bitfield.set(bitfield_index, true);
aggregation_bitfield.set(committee_size, false);
let mut custody_bitfield = Bitfield::new();
@ -316,26 +339,20 @@ where
agg_sig.add(&Signature::new(
&message,
domain,
self.get_sk(*validator_index),
self.get_sk(validator_index),
));
agg_sig
};
let attestation = Attestation {
Attestation {
aggregation_bitfield,
data,
custody_bitfield,
signature,
};
}
}
self.chain
.process_attestation(attestation)
.expect("should process attestation");
}
}
});
}
/// Returns the secret key for the given validator index.
fn get_sk(&self, validator_index: usize) -> &SecretKey {