Move BeaconStateBuilder out of test mods

This allows it to be accessed by other crates/tests.
This commit is contained in:
Paul Hauner 2019-02-26 16:25:08 +13:00
parent 87ea95ce35
commit 59d6b1fdd0
No known key found for this signature in database
GPG Key ID: 303E4494BB28068C
3 changed files with 221 additions and 64 deletions

View File

@ -14,6 +14,9 @@ use ssz::{hash, Decodable, DecodeError, Encodable, SszStream, TreeHash};
use std::collections::HashMap;
use swap_or_not_shuffle::get_permutated_index;
pub use builder::BeaconStateBuilder;
mod builder;
mod epoch_cache;
mod tests;

View File

@ -0,0 +1,215 @@
use crate::*;
use bls::create_proof_of_possession;
/// Builds a `BeaconState` for use in testing or benchmarking.
///
/// Building the `BeaconState` is a three step processes:
///
/// 1. Create a new `BeaconStateBuilder`.
/// 2. Run the `genesis` function to generate a new BeaconState.
/// 3. (Optional) Use builder functions to modify the `BeaconState`.
/// 4. Call `build()` to obtain a cloned `BeaconState`.
///
/// Step (2) is necessary because step (3) requires an existing `BeaconState` object. (2) is not
/// included in (1) to allow for modifying params before generating the `BeaconState` (e.g., the
/// spec).
///
/// Step (4) produces a clone of the BeaconState and doesn't consume the `BeaconStateBuilder` to
/// allow access to the `keypairs` and `spec`.
pub struct BeaconStateBuilder {
pub state: Option<BeaconState>,
pub genesis_time: u64,
pub initial_validator_deposits: Vec<Deposit>,
pub latest_eth1_data: Eth1Data,
pub spec: ChainSpec,
pub keypairs: Vec<Keypair>,
}
impl BeaconStateBuilder {
/// Create a new builder with the given number of validators.
pub fn with_random_validators(validator_count: usize) -> Self {
let genesis_time = 10_000_000;
let keypairs: Vec<Keypair> = (0..validator_count)
.collect::<Vec<usize>>()
.iter()
.map(|_| Keypair::random())
.collect();
let initial_validator_deposits = keypairs
.iter()
.map(|keypair| Deposit {
branch: vec![], // branch verification is not specified.
index: 0, // index verification is not specified.
deposit_data: DepositData {
amount: 32_000_000_000, // 32 ETH (in Gwei)
timestamp: genesis_time - 1,
deposit_input: DepositInput {
pubkey: keypair.pk.clone(),
withdrawal_credentials: Hash256::zero(), // Withdrawal not possible.
proof_of_possession: create_proof_of_possession(&keypair),
},
},
})
.collect();
let latest_eth1_data = Eth1Data {
deposit_root: Hash256::zero(),
block_hash: Hash256::zero(),
};
let spec = ChainSpec::foundation();
Self {
state: None,
genesis_time,
initial_validator_deposits,
latest_eth1_data,
spec,
keypairs,
}
}
/// Runs the `BeaconState::genesis` function and produces a `BeaconState`.
pub fn genesis(&mut self) -> Result<(), BeaconStateError> {
let state = BeaconState::genesis(
self.genesis_time,
self.initial_validator_deposits.clone(),
self.latest_eth1_data.clone(),
&self.spec,
)?;
self.state = Some(state);
Ok(())
}
/// Sets the `BeaconState` to be in the last slot of the given epoch.
///
/// Sets all justification/finalization parameters to be be as "perfect" as possible (i.e.,
/// highest justified and finalized slots, full justification bitfield, etc).
pub fn teleport_to_end_of_epoch(&mut self, epoch: Epoch) {
let state = self.state.as_mut().expect("Genesis required");
let slot = epoch.end_slot(self.spec.epoch_length);
state.slot = slot;
state.validator_registry_update_epoch = epoch - 1;
state.previous_calculation_epoch = epoch - 1;
state.current_calculation_epoch = epoch;
state.previous_epoch_seed = Hash256::from(&b"previous_seed"[..]);
state.current_epoch_seed = Hash256::from(&b"current_seed"[..]);
state.previous_justified_epoch = epoch - 2;
state.justified_epoch = epoch - 1;
state.justification_bitfield = u64::max_value();
state.finalized_epoch = epoch - 1;
}
/// Creates a full set of attestations for the `BeaconState`. Each attestation has full
/// participation from its committee and references the expected beacon_block hashes.
///
/// These attestations should be fully conducive to justification and finalization.
pub fn insert_attestations(&mut self) {
let state = self.state.as_mut().expect("Genesis required");
state
.build_epoch_cache(RelativeEpoch::Previous, &self.spec)
.unwrap();
state
.build_epoch_cache(RelativeEpoch::Current, &self.spec)
.unwrap();
let current_epoch = state.current_epoch(&self.spec);
let previous_epoch = state.previous_epoch(&self.spec);
let current_epoch_depth =
(state.slot - current_epoch.end_slot(self.spec.epoch_length)).as_usize();
let previous_epoch_slots = previous_epoch.slot_iter(self.spec.epoch_length);
let current_epoch_slots = current_epoch
.slot_iter(self.spec.epoch_length)
.take(current_epoch_depth);
for slot in previous_epoch_slots.chain(current_epoch_slots) {
let committees = state
.get_crosslink_committees_at_slot(slot, &self.spec)
.unwrap()
.clone();
for (committee, shard) in committees {
state
.latest_attestations
.push(committee_to_pending_attestation(
state, &committee, shard, slot, &self.spec,
))
}
}
}
/// Returns a cloned `BeaconState`.
pub fn build(&self) -> Result<BeaconState, BeaconStateError> {
match &self.state {
Some(state) => Ok(state.clone()),
None => panic!("Genesis required"),
}
}
}
/// Builds a valid PendingAttestation with full participation for some committee.
fn committee_to_pending_attestation(
state: &BeaconState,
committee: &Vec<usize>,
shard: u64,
slot: Slot,
spec: &ChainSpec,
) -> PendingAttestation {
let current_epoch = state.current_epoch(spec);
let previous_epoch = state.previous_epoch(spec);
let mut aggregation_bitfield = Bitfield::new();
let mut custody_bitfield = Bitfield::new();
for (i, _) in committee.iter().enumerate() {
aggregation_bitfield.set(i, true);
custody_bitfield.set(i, true);
}
let is_previous_epoch = state.slot.epoch(spec.epoch_length) != slot.epoch(spec.epoch_length);
let justified_epoch = if is_previous_epoch {
state.previous_justified_epoch
} else {
state.justified_epoch
};
let epoch_boundary_root = if is_previous_epoch {
*state
.get_block_root(previous_epoch.start_slot(spec.epoch_length), spec)
.unwrap()
} else {
*state
.get_block_root(current_epoch.start_slot(spec.epoch_length), spec)
.unwrap()
};
let justified_block_root = *state
.get_block_root(justified_epoch.start_slot(spec.epoch_length), &spec)
.unwrap();
PendingAttestation {
aggregation_bitfield,
data: AttestationData {
slot,
shard,
beacon_block_root: *state.get_block_root(slot, spec).unwrap(),
epoch_boundary_root,
shard_block_root: Hash256::zero(),
latest_crosslink: Crosslink {
epoch: slot.epoch(spec.epoch_length),
shard_block_root: Hash256::zero(),
},
justified_epoch,
justified_block_root,
},
custody_bitfield,
inclusion_slot: slot,
}
}

View File

@ -2,73 +2,12 @@
use super::*;
use crate::test_utils::{SeedableRng, TestRandom, XorShiftRng};
use crate::{
BeaconState, BeaconStateError, ChainSpec, Deposit, DepositData, DepositInput, Eth1Data,
Hash256, Keypair,
};
use bls::create_proof_of_possession;
use crate::{BeaconState, ChainSpec};
use ssz::{ssz_encode, Decodable};
struct BeaconStateTestBuilder {
pub genesis_time: u64,
pub initial_validator_deposits: Vec<Deposit>,
pub latest_eth1_data: Eth1Data,
pub spec: ChainSpec,
pub keypairs: Vec<Keypair>,
}
impl BeaconStateTestBuilder {
pub fn with_random_validators(validator_count: usize) -> Self {
let genesis_time = 10_000_000;
let keypairs: Vec<Keypair> = (0..validator_count)
.collect::<Vec<usize>>()
.iter()
.map(|_| Keypair::random())
.collect();
let initial_validator_deposits = keypairs
.iter()
.map(|keypair| Deposit {
branch: vec![], // branch verification is not specified.
index: 0, // index verification is not specified.
deposit_data: DepositData {
amount: 32_000_000_000, // 32 ETH (in Gwei)
timestamp: genesis_time - 1,
deposit_input: DepositInput {
pubkey: keypair.pk.clone(),
withdrawal_credentials: Hash256::zero(), // Withdrawal not possible.
proof_of_possession: create_proof_of_possession(&keypair),
},
},
})
.collect();
let latest_eth1_data = Eth1Data {
deposit_root: Hash256::zero(),
block_hash: Hash256::zero(),
};
let spec = ChainSpec::foundation();
Self {
genesis_time,
initial_validator_deposits,
latest_eth1_data,
spec,
keypairs,
}
}
pub fn build(&self) -> Result<BeaconState, BeaconStateError> {
BeaconState::genesis(
self.genesis_time,
self.initial_validator_deposits.clone(),
self.latest_eth1_data.clone(),
&self.spec,
)
}
}
#[test]
pub fn can_produce_genesis_block() {
let builder = BeaconStateTestBuilder::with_random_validators(2);
let builder = BeaconStateBuilder::with_random_validators(2);
builder.build().unwrap();
}
@ -79,7 +18,7 @@ pub fn can_produce_genesis_block() {
pub fn get_attestation_participants_consistency() {
let mut rng = XorShiftRng::from_seed([42; 16]);
let mut builder = BeaconStateTestBuilder::with_random_validators(8);
let mut builder = BeaconStateBuilder::with_random_validators(8);
builder.spec = ChainSpec::few_validators();
let mut state = builder.build().unwrap();