Add benches, examples to fork_choice

eth2/fork_choice/Cargo.toml

@@ -4,6 +4,10 @@ version = "0.1.0"
 authors = ["Age Manning <Age@AgeManning.com>"]
 edition = "2018"
 
+[[bench]]
+name = "benches"
+harness = false
+
 [dependencies]
 store = { path = "../../beacon_node/store" }
 ssz = { path = "../utils/ssz" }
@@ -12,6 +16,7 @@ log = "0.4.6"
 bit-vec = "0.5.0"
 
 [dev-dependencies]
+criterion = "0.2"
 hex = "0.3.2"
 yaml-rust = "0.4.2"
 bls = { path = "../utils/bls" }

eth2/fork_choice/benches/benches.rs

@@ -0,0 +1,75 @@
+use criterion::Criterion;
+use criterion::{criterion_group, criterion_main, Benchmark};
+use fork_choice::{test_utils::TestingForkChoiceBuilder, ForkChoice, OptimizedLMDGhost};
+use std::sync::Arc;
+use store::MemoryStore;
+use types::{ChainSpec, EthSpec, FoundationEthSpec};
+
+pub type TestedForkChoice<T, U> = OptimizedLMDGhost<T, U>;
+pub type TestedEthSpec = FoundationEthSpec;
+
+/// Helper function to setup a builder and spec.
+fn setup(
+    validator_count: usize,
+    chain_length: usize,
+) -> (
+    TestingForkChoiceBuilder<MemoryStore, TestedEthSpec>,
+    ChainSpec,
+) {
+    let store = MemoryStore::open();
+    let builder: TestingForkChoiceBuilder<MemoryStore, TestedEthSpec> =
+        TestingForkChoiceBuilder::new(validator_count, chain_length, Arc::new(store));
+    let spec = TestedEthSpec::spec();
+
+    (builder, spec)
+}
+
+/// Benches adding blocks to fork_choice.
+fn add_block(c: &mut Criterion) {
+    let validator_count = 16;
+    let chain_length = 100;
+
+    let (builder, spec) = setup(validator_count, chain_length);
+
+    c.bench(
+        &format!("{}_blocks", chain_length),
+        Benchmark::new("add_blocks", move |b| {
+            b.iter(|| {
+                let mut fc = builder.build::<TestedForkChoice<MemoryStore, TestedEthSpec>>();
+                for (root, block) in builder.chain.iter().skip(1) {
+                    fc.add_block(block, root, &spec).unwrap();
+                }
+            })
+        })
+        .sample_size(10),
+    );
+}
+
+/// Benches fork choice head finding.
+fn find_head(c: &mut Criterion) {
+    let validator_count = 16;
+    let chain_length = 64 * 2;
+
+    let (builder, spec) = setup(validator_count, chain_length);
+
+    let mut fc = builder.build::<TestedForkChoice<MemoryStore, TestedEthSpec>>();
+    for (root, block) in builder.chain.iter().skip(1) {
+        fc.add_block(block, root, &spec).unwrap();
+    }
+
+    let head_root = builder.chain.last().unwrap().0;
+    for i in 0..validator_count {
+        fc.add_attestation(i as u64, &head_root, &spec).unwrap();
+    }
+
+    c.bench(
+        &format!("{}_blocks", chain_length),
+        Benchmark::new("find_head", move |b| {
+            b.iter(|| fc.find_head(&builder.genesis_root(), &spec).unwrap())
+        })
+        .sample_size(10),
+    );
+}
+
+criterion_group!(benches, add_block, find_head);
+criterion_main!(benches);

eth2/fork_choice/examples/example.rs

@@ -0,0 +1,40 @@
+use fork_choice::{test_utils::TestingForkChoiceBuilder, ForkChoice, OptimizedLMDGhost};
+use std::sync::Arc;
+use store::{MemoryStore, Store};
+use types::{BeaconBlock, ChainSpec, EthSpec, FoundationEthSpec, Hash256};
+
+fn main() {
+    let validator_count = 16;
+    let chain_length = 100;
+    let repetitions = 50;
+
+    let store = MemoryStore::open();
+    let builder: TestingForkChoiceBuilder<MemoryStore, FoundationEthSpec> =
+        TestingForkChoiceBuilder::new(validator_count, chain_length, Arc::new(store));
+
+    let fork_choosers: Vec<OptimizedLMDGhost<MemoryStore, FoundationEthSpec>> = (0..repetitions)
+        .into_iter()
+        .map(|_| builder.build())
+        .collect();
+
+    let spec = &FoundationEthSpec::spec();
+
+    println!("Running {} times...", repetitions);
+    for fc in fork_choosers {
+        do_thing(fc, &builder.chain, builder.genesis_root(), spec);
+    }
+}
+
+#[inline(never)]
+fn do_thing<F: ForkChoice<S>, S: Store>(
+    mut fc: F,
+    chain: &[(Hash256, BeaconBlock)],
+    genesis_root: Hash256,
+    spec: &ChainSpec,
+) {
+    for (root, block) in chain.iter().skip(1) {
+        fc.add_block(block, root, spec).unwrap();
+    }
+
+    let _head = fc.find_head(&genesis_root, spec).unwrap();
+}

eth2/fork_choice/src/lib.rs

@@ -20,6 +20,7 @@ pub mod bitwise_lmd_ghost;
 pub mod longest_chain;
 pub mod optimized_lmd_ghost;
 pub mod slow_lmd_ghost;
+pub mod test_utils;
 
 use std::sync::Arc;
 use store::Error as DBError;

eth2/fork_choice/src/optimized_lmd_ghost.rs

@@ -69,12 +69,11 @@ impl<T: Store, E: EthSpec> OptimizedLMDGhost<T, E> {
 
         let active_validator_indices =
             current_state.get_active_validator_indices(block_slot.epoch(spec.slots_per_epoch));
+        let validator_balances = &current_state.validator_balances;
 
         for index in active_validator_indices {
-            let balance = std::cmp::min(
-                current_state.validator_balances[index],
-                spec.max_deposit_amount,
-            ) / spec.fork_choice_balance_increment;
+            let balance = std::cmp::min(validator_balances[index], spec.max_deposit_amount)
+                / spec.fork_choice_balance_increment;
             if balance > 0 {
                 if let Some(target) = self.latest_attestation_targets.get(&(index as u64)) {
                     *latest_votes.entry(*target).or_insert_with(|| 0) += balance;

eth2/fork_choice/src/test_utils.rs

@@ -0,0 +1,94 @@
+use crate::ForkChoice;
+use std::marker::PhantomData;
+use std::sync::Arc;
+use store::Store;
+use types::{
+    test_utils::{SeedableRng, TestRandom, TestingBeaconStateBuilder, XorShiftRng},
+    BeaconBlock, BeaconState, EthSpec, FoundationEthSpec, Hash256, Keypair,
+};
+
+/// Creates a chain of blocks and produces `ForkChoice` instances with pre-filled stores.
+pub struct TestingForkChoiceBuilder<S, E> {
+    store: Arc<S>,
+    pub chain: Vec<(Hash256, BeaconBlock)>,
+    _phantom: PhantomData<E>,
+}
+
+impl<S: Store, E: EthSpec> TestingForkChoiceBuilder<S, E> {
+    pub fn new(validator_count: usize, chain_length: usize, store: Arc<S>) -> Self {
+        let chain = get_chain_of_blocks::<FoundationEthSpec, S>(
+            chain_length,
+            validator_count,
+            store.clone(),
+        );
+
+        Self {
+            store,
+            chain,
+            _phantom: PhantomData,
+        }
+    }
+
+    pub fn genesis_root(&self) -> Hash256 {
+        self.chain[0].0
+    }
+
+    /// Return a new `ForkChoice` instance with a chain stored in it's `Store`.
+    pub fn build<F: ForkChoice<S>>(&self) -> F {
+        F::new(self.store.clone())
+    }
+}
+
+fn get_state<T: EthSpec>(validator_count: usize) -> BeaconState<T> {
+    let spec = &T::spec();
+
+    let builder: TestingBeaconStateBuilder<T> =
+        TestingBeaconStateBuilder::from_single_keypair(validator_count, &Keypair::random(), spec);
+    let (state, _keypairs) = builder.build();
+    state
+}
+
+/// Generates a chain of blocks of length `len`.
+///
+/// Creates a `BeaconState` for the block and stores it in `Store`, along with the block.
+///
+/// Returns the chain of blocks.
+fn get_chain_of_blocks<T: EthSpec, U: Store>(
+    len: usize,
+    validator_count: usize,
+    store: Arc<U>,
+) -> Vec<(Hash256, BeaconBlock)> {
+    let spec = T::spec();
+    let mut blocks_and_roots: Vec<(Hash256, BeaconBlock)> = vec![];
+    let mut unique_hashes = (0..).into_iter().map(|i| Hash256::from(i));
+    let mut random_block = BeaconBlock::random_for_test(&mut XorShiftRng::from_seed([42; 16]));
+    random_block.previous_block_root = Hash256::zero();
+    let beacon_state = get_state::<T>(validator_count);
+
+    for i in 0..len {
+        let slot = spec.genesis_slot + i as u64;
+
+        // Generate and store the state.
+        let mut state = beacon_state.clone();
+        state.slot = slot;
+        let state_root = unique_hashes.next().unwrap();
+        store.put(&state_root, &state).unwrap();
+
+        // Generate the block.
+        let mut block = random_block.clone();
+        block.slot = slot;
+        block.state_root = state_root;
+
+        // Chain all the blocks to their parents.
+        if i > 0 {
+            block.previous_block_root = blocks_and_roots[i - 1].0;
+        }
+
+        // Store the block.
+        let block_root = unique_hashes.next().unwrap();
+        store.put(&block_root, &block).unwrap();
+        blocks_and_roots.push((block_root, block));
+    }
+
+    blocks_and_roots
+}