Generalise fork choice tests.

beacon_node/beacon_chain/src/lib.rs

@@ -5,4 +5,4 @@ mod checkpoint;
 
 pub use self::beacon_chain::{BeaconChain, Error};
 pub use self::checkpoint::CheckPoint;
-pub use fork_choice::{ForkChoice, ForkChoiceAlgorithms, ForkChoiceError};
+pub use fork_choice::{ForkChoice, ForkChoiceAlgorithm, ForkChoiceError};

eth2/fork_choice/src/lib.rs

@@ -1,22 +1,20 @@
 //! This crate stores the various implementations of fork-choice rules that can be used for the
 //! beacon blockchain.
 //!
-//! There are four implementations. One is the naive longest chain rule (primarily for testing
+//! There are three implementations. One is the naive longest chain rule (primarily for testing
-//! purposes). The other three are proposed implementations of the LMD-GHOST fork-choice rule with various forms of optimisation.
+//! purposes). The other two are proposed implementations of the LMD-GHOST fork-choice rule with various forms of optimisation.
 //!
 //! The current implementations are:
 //! - [`longest-chain`]: Simplistic longest-chain fork choice - primarily for testing, **not for
 //! production**.
 //! - [`slow_lmd_ghost`]: This is a simple and very inefficient implementation given in the ethereum 2.0
 //! specifications (https://github.com/ethereum/eth2.0-specs/blob/v0.1/specs/core/0_beacon-chain.md#get_block_root).
-//! - [`optimised_lmd_ghost`]: This is an optimised version of the naive implementation as proposed
+//! - [`optimised_lmd_ghost`]: This is an optimised version of bitwise LMD-GHOST as proposed
 //! by Vitalik. The reference implementation can be found at: https://github.com/ethereum/research/blob/master/ghost/ghost.py
-//! - [`protolambda_lmd_ghost`]: Another optimised version of LMD-GHOST designed by @protolambda.
-//! The go implementation can be found here: https://github.com/protolambda/lmd-ghost.
 //!
+//! [`longest-chain`]: struct.LongestChain.html
 //! [`slow_lmd_ghost`]: struct.SlowLmdGhost.html
 //! [`optimised_lmd_ghost`]: struct.OptimisedLmdGhost.html
-//! [`protolambda_lmd_ghost`]: struct.ProtolambdaLmdGhost.html
 
 extern crate db;
 extern crate ssz;
@@ -34,6 +32,7 @@ use types::{BeaconBlock, ChainSpec, Hash256};
 
 pub use longest_chain::LongestChain;
 pub use optimised_lmd_ghost::OptimisedLMDGhost;
+pub use slow_lmd_ghost::SlowLMDGhost;
 
 /// Defines the interface for Fork Choices. Each Fork choice will define their own data structures
 /// which can be built in block processing through the `add_block` and `add_attestation` functions.
@@ -97,7 +96,7 @@ impl From<BeaconBlockAtSlotError> for ForkChoiceError {
 }
 
 /// Fork choice options that are currently implemented.
-pub enum ForkChoiceAlgorithms {
+pub enum ForkChoiceAlgorithm {
     /// Chooses the longest chain becomes the head. Not for production.
     LongestChain,
     /// A simple and highly inefficient implementation of LMD ghost.

eth2/fork_choice/src/slow_lmd_ghost.rs

@@ -30,12 +30,15 @@ impl<T> SlowLMDGhost<T>
 where
     T: ClientDB + Sized,
 {
-    pub fn new(block_store: BeaconBlockStore<T>, state_store: BeaconStateStore<T>) -> Self {
+    pub fn new(
+        block_store: Arc<BeaconBlockStore<T>>,
+        state_store: Arc<BeaconStateStore<T>>,
+    ) -> Self {
         SlowLMDGhost {
             latest_attestation_targets: HashMap::new(),
             children: HashMap::new(),
-            block_store: Arc::new(block_store),
+            block_store,
-            state_store: Arc::new(state_store),
+            state_store,
         }
     }
 

eth2/fork_choice/tests/tests.rs

@@ -15,7 +15,7 @@ use bls::{PublicKey, Signature};
 use db::stores::{BeaconBlockStore, BeaconStateStore};
 use db::MemoryDB;
 use env_logger::{Builder, Env};
-use fork_choice::{ForkChoice, OptimisedLMDGhost};
+use fork_choice::{ForkChoice, ForkChoiceAlgorithm, LongestChain, OptimisedLMDGhost, SlowLMDGhost};
 use ssz::ssz_encode;
 use std::collections::HashMap;
 use std::sync::Arc;
@@ -25,86 +25,32 @@ use types::{
 };
 use yaml_rust::yaml;
 
-// initialise a single validator and state. All blocks will reference this state root.
+// run tests
-fn setup_inital_state(
-    no_validators: usize,
-) -> (impl ForkChoice, Arc<BeaconBlockStore<MemoryDB>>, Hash256) {
-    let zero_hash = Hash256::zero();
-
-    let db = Arc::new(MemoryDB::open());
-    let block_store = Arc::new(BeaconBlockStore::new(db.clone()));
-    let state_store = Arc::new(BeaconStateStore::new(db.clone()));
-
-    // the fork choice instantiation
-    let optimised_lmd_ghost = OptimisedLMDGhost::new(block_store.clone(), state_store.clone());
-
-    // misc vars for setting up the state
-    let genesis_time = 1_550_381_159;
-
-    let latest_eth1_data = Eth1Data {
-        deposit_root: zero_hash.clone(),
-        block_hash: zero_hash.clone(),
-    };
-
-    let initial_validator_deposits = vec![];
-    let spec = ChainSpec::foundation();
-
-    // create the state
-    let mut state = BeaconState::genesis(
-        genesis_time,
-        initial_validator_deposits,
-        latest_eth1_data,
-        &spec,
-    )
-    .unwrap();
-
-    let default_validator = Validator {
-        pubkey: PublicKey::default(),
-        withdrawal_credentials: zero_hash,
-        activation_epoch: Epoch::from(0u64),
-        exit_epoch: spec.far_future_epoch,
-        withdrawal_epoch: spec.far_future_epoch,
-        penalized_epoch: spec.far_future_epoch,
-        status_flags: None,
-    };
-    // activate the validators
-    for _ in 0..no_validators {
-        state.validator_registry.push(default_validator.clone());
-        state.validator_balances.push(32_000_000_000);
-    }
-
-    let state_root = state.canonical_root();
-    state_store
-        .put(&state_root, &ssz_encode(&state)[..])
-        .unwrap();
-
-    // return initialised vars
-    (optimised_lmd_ghost, block_store, state_root)
-}
-
-// YAML test vectors
 #[test]
 fn test_optimised_lmd_ghost() {
+    test_yaml_vectors(
+        ForkChoiceAlgorithm::OptimisedLMDGhost,
+        "tests/optimised_lmd_ghost_test_vectors.yaml",
+        100,
+    );
+}
+
+// run a generic test over given YAML test vectors
+fn test_yaml_vectors(
+    fork_choice_algo: ForkChoiceAlgorithm,
+    yaml_file_path: &str,
+    max_validators: usize,
+) {
     // set up logging
     Builder::from_env(Env::default().default_filter_or("debug")).init();
 
-    // load the yaml
+    // load test cases from yaml
-    let mut file = {
+    let test_cases = load_test_cases_from_yaml(yaml_file_path);
-        let mut file_path_buf = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
-        file_path_buf.push("tests/fork_choice_test_vectors.yaml");
-
-        File::open(file_path_buf).unwrap()
-    };
-
-    let mut yaml_str = String::new();
-    file.read_to_string(&mut yaml_str).unwrap();
-    let docs = yaml::YamlLoader::load_from_str(&yaml_str).unwrap();
-    let doc = &docs[0];
-    let test_cases = doc["test_cases"].as_vec().unwrap();
 
     // set up the test
-    let total_emulated_validators = 20; // the number of validators used to give weights.
+    let total_emulated_validators = max_validators; // the number of validators used to give weights.
-    let (mut fork_choice, block_store, state_root) = setup_inital_state(total_emulated_validators);
+    let (mut fork_choice, block_store, state_root) =
+        setup_inital_state(fork_choice_algo, total_emulated_validators);
 
     // keep a hashmap of block_id's to block_hashes (random hashes to abstract block_id)
     let mut block_id_map: HashMap<String, Hash256> = HashMap::new();
@@ -227,3 +173,85 @@ fn test_optimised_lmd_ghost() {
         assert!(success, "Did not find one of the possible heads");
     }
 }
+
+// loads the test_cases from the supplied yaml file
+fn load_test_cases_from_yaml(file_path: &str) -> Vec<yaml_rust::Yaml> {
+    // load the yaml
+    let mut file = {
+        let mut file_path_buf = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+        file_path_buf.push(file_path);
+        File::open(file_path_buf).unwrap()
+    };
+    let mut yaml_str = String::new();
+    file.read_to_string(&mut yaml_str).unwrap();
+    let docs = yaml::YamlLoader::load_from_str(&yaml_str).unwrap();
+    let doc = &docs[0];
+    doc["test_cases"].as_vec().unwrap().clone()
+}
+
+// initialise a single validator and state. All blocks will reference this state root.
+fn setup_inital_state(
+    fork_choice_algo: ForkChoiceAlgorithm,
+    no_validators: usize,
+) -> (Box<ForkChoice>, Arc<BeaconBlockStore<MemoryDB>>, Hash256) {
+    let zero_hash = Hash256::zero();
+
+    let db = Arc::new(MemoryDB::open());
+    let block_store = Arc::new(BeaconBlockStore::new(db.clone()));
+    let state_store = Arc::new(BeaconStateStore::new(db.clone()));
+
+    // the fork choice instantiation
+    let fork_choice: Box<ForkChoice> = match fork_choice_algo {
+        ForkChoiceAlgorithm::OptimisedLMDGhost => Box::new(OptimisedLMDGhost::new(
+            block_store.clone(),
+            state_store.clone(),
+        )),
+        ForkChoiceAlgorithm::SlowLMDGhost => {
+            Box::new(SlowLMDGhost::new(block_store.clone(), state_store.clone()))
+        }
+        ForkChoiceAlgorithm::LongestChain => Box::new(LongestChain::new(block_store.clone())),
+    };
+
+    // misc vars for setting up the state
+    let genesis_time = 1_550_381_159;
+
+    let latest_eth1_data = Eth1Data {
+        deposit_root: zero_hash.clone(),
+        block_hash: zero_hash.clone(),
+    };
+
+    let initial_validator_deposits = vec![];
+    let spec = ChainSpec::foundation();
+
+    // create the state
+    let mut state = BeaconState::genesis(
+        genesis_time,
+        initial_validator_deposits,
+        latest_eth1_data,
+        &spec,
+    )
+    .unwrap();
+
+    let default_validator = Validator {
+        pubkey: PublicKey::default(),
+        withdrawal_credentials: zero_hash,
+        activation_epoch: Epoch::from(0u64),
+        exit_epoch: spec.far_future_epoch,
+        withdrawal_epoch: spec.far_future_epoch,
+        penalized_epoch: spec.far_future_epoch,
+        status_flags: None,
+    };
+    // activate the validators
+    for _ in 0..no_validators {
+        state.validator_registry.push(default_validator.clone());
+        state.validator_balances.push(32_000_000_000);
+    }
+
+    let state_root = state.canonical_root();
+    state_store
+        .put(&state_root, &ssz_encode(&state)[..])
+        .unwrap();
+
+    // return initialised vars
+    (fork_choice, block_store, state_root)
+}