Merge pull request #195 from sigp/fork-choices

[WIP] Generalise Fork Choice Algorithms and implement optimised version.

Cargo.toml

@@ -2,7 +2,7 @@
 members = [
 	"eth2/attester",
 	"eth2/block_producer",
-	"eth2/naive_fork_choice",
+	"eth2/fork_choice",
 	"eth2/state_processing",
 	"eth2/types",
 	"eth2/utils/bls",

beacon_node/Cargo.toml

@@ -14,6 +14,7 @@ clap = "2.32.0"
 db = { path = "db" }
 dirs = "1.0.3"
 futures = "0.1.23"
+fork_choice = { path = "../eth2/fork_choice" }
 slog = "^2.2.3"
 slot_clock = { path = "../eth2/utils/slot_clock" }
 slog-term = "^2.4.0"

beacon_node/beacon_chain/Cargo.toml

@@ -12,6 +12,7 @@ db = { path = "../db" }
 failure = "0.1"
 failure_derive = "0.1"
 hashing = { path = "../../eth2/utils/hashing" }
+fork_choice = { path = "../../eth2/fork_choice" }
 parking_lot = "0.7"
 log = "0.4"
 env_logger = "0.6"

beacon_node/beacon_chain/src/attestation_targets.rs

@@ -1,23 +0,0 @@
-use std::collections::HashMap;
-use types::Hash256;
-
-#[derive(Default)]
-pub struct AttestationTargets {
-    map: HashMap<u64, Hash256>,
-}
-
-impl AttestationTargets {
-    pub fn new() -> Self {
-        Self {
-            map: HashMap::new(),
-        }
-    }
-
-    pub fn get(&self, validator_index: u64) -> Option<&Hash256> {
-        self.map.get(&validator_index)
-    }
-
-    pub fn insert(&mut self, validator_index: u64, block_hash: Hash256) -> Option<Hash256> {
-        self.map.insert(validator_index, block_hash)
-    }
-}

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -1,11 +1,10 @@
 use crate::attestation_aggregator::{AttestationAggregator, Outcome as AggregationOutcome};
-use crate::attestation_targets::AttestationTargets;
-use crate::block_graph::BlockGraph;
 use crate::checkpoint::CheckPoint;
 use db::{
     stores::{BeaconBlockStore, BeaconStateStore},
     ClientDB, DBError,
 };
+use fork_choice::{ForkChoice, ForkChoiceError};
 use log::{debug, trace};
 use parking_lot::{RwLock, RwLockReadGuard};
 use slot_clock::SlotClock;
@@ -28,11 +27,14 @@ pub enum Error {
     CommitteesError(CommitteesError),
     DBInconsistent(String),
     DBError(String),
+    ForkChoiceError(ForkChoiceError),
+    MissingBeaconBlock(Hash256),
+    MissingBeaconState(Hash256),
 }
 
 #[derive(Debug, PartialEq)]
 pub enum ValidBlock {
-    /// The block was sucessfully processed.
+    /// The block was successfully processed.
     Processed,
 }
 
@@ -53,29 +55,29 @@ pub enum InvalidBlock {
 
 #[derive(Debug, PartialEq)]
 pub enum BlockProcessingOutcome {
-    /// The block was sucessfully validated.
+    /// The block was successfully validated.
     ValidBlock(ValidBlock),
-    /// The block was not sucessfully validated.
+    /// The block was not successfully validated.
     InvalidBlock(InvalidBlock),
 }
 
-pub struct BeaconChain<T: ClientDB + Sized, U: SlotClock> {
+pub struct BeaconChain<T: ClientDB + Sized, U: SlotClock, F: ForkChoice> {
     pub block_store: Arc<BeaconBlockStore<T>>,
     pub state_store: Arc<BeaconStateStore<T>>,
     pub slot_clock: U,
-    pub block_graph: BlockGraph,
     pub attestation_aggregator: RwLock<AttestationAggregator>,
     canonical_head: RwLock<CheckPoint>,
     finalized_head: RwLock<CheckPoint>,
     pub state: RwLock<BeaconState>,
-    pub latest_attestation_targets: RwLock<AttestationTargets>,
     pub spec: ChainSpec,
+    pub fork_choice: RwLock<F>,
 }
 
-impl<T, U> BeaconChain<T, U>
+impl<T, U, F> BeaconChain<T, U, F>
 where
     T: ClientDB,
     U: SlotClock,
+    F: ForkChoice,
 {
     /// Instantiate a new Beacon Chain, from genesis.
     pub fn genesis(
@@ -86,6 +88,7 @@ where
         latest_eth1_data: Eth1Data,
         initial_validator_deposits: Vec<Deposit>,
         spec: ChainSpec,
+        fork_choice: F,
     ) -> Result<Self, Error> {
         if initial_validator_deposits.is_empty() {
             return Err(Error::InsufficientValidators);
@@ -104,9 +107,6 @@ where
         let block_root = genesis_block.canonical_root();
         block_store.put(&block_root, &ssz_encode(&genesis_block)[..])?;
 
-        let block_graph = BlockGraph::new();
-        block_graph.add_leaf(&Hash256::zero(), block_root);
-
         let finalized_head = RwLock::new(CheckPoint::new(
             genesis_block.clone(),
             block_root,
@@ -121,19 +121,16 @@ where
         ));
         let attestation_aggregator = RwLock::new(AttestationAggregator::new());
 
-        let latest_attestation_targets = RwLock::new(AttestationTargets::new());
-
         Ok(Self {
             block_store,
             state_store,
             slot_clock,
-            block_graph,
             attestation_aggregator,
             state: RwLock::new(genesis_state.clone()),
             finalized_head,
             canonical_head,
-            latest_attestation_targets,
             spec,
+            fork_choice: RwLock::new(fork_choice),
         })
     }
 
@@ -208,7 +205,7 @@ where
         Ok(())
     }
 
-    /// Returns the the validator index (if any) for the given public key.
+    /// Returns the validator index (if any) for the given public key.
     ///
     /// Information is retrieved from the present `beacon_state.validator_registry`.
     pub fn validator_index(&self, pubkey: &PublicKey) -> Option<usize> {
@@ -332,26 +329,24 @@ where
         &self,
         free_attestation: FreeAttestation,
     ) -> Result<AggregationOutcome, Error> {
-        self.attestation_aggregator
+        let aggregation_outcome = self
+            .attestation_aggregator
             .write()
-            .process_free_attestation(&self.state.read(), &free_attestation, &self.spec)
+            .process_free_attestation(&self.state.read(), &free_attestation, &self.spec)?;
-            .map_err(|e| e.into())
+        // TODO: Check this comment
-    }
+        //.map_err(|e| e.into())?;
 
-    /// Set the latest attestation target for some validator.
+        // return if the attestation is invalid
-    pub fn insert_latest_attestation_target(&self, validator_index: u64, block_root: Hash256) {
+        if !aggregation_outcome.valid {
-        let mut targets = self.latest_attestation_targets.write();
+            return Ok(aggregation_outcome);
-        targets.insert(validator_index, block_root);
-    }
-
-    /// Get the latest attestation target for some validator.
-    pub fn get_latest_attestation_target(&self, validator_index: u64) -> Option<Hash256> {
-        let targets = self.latest_attestation_targets.read();
-
-        match targets.get(validator_index) {
-            Some(hash) => Some(*hash),
-            None => None,
         }
+
+        // valid attestation, proceed with fork-choice logic
+        self.fork_choice.write().add_attestation(
+            free_attestation.validator_index,
+            &free_attestation.data.beacon_block_root,
+        )?;
+        Ok(aggregation_outcome)
     }
 
     /// Dumps the entire canonical chain, from the head to genesis to a vector for analysis.
@@ -458,7 +453,7 @@ where
             }
         }
 
-        // Apply the recieved block to its parent state (which has been transitioned into this
+        // Apply the received block to its parent state (which has been transitioned into this
         // slot).
         if let Err(e) = state.per_block_processing(&block, &self.spec) {
             return Ok(BlockProcessingOutcome::InvalidBlock(
@@ -478,15 +473,20 @@ where
         self.block_store.put(&block_root, &ssz_encode(&block)[..])?;
         self.state_store.put(&state_root, &ssz_encode(&state)[..])?;
 
-        // Update the block DAG.
+        // run the fork_choice add_block logic
-        self.block_graph.add_leaf(&parent_block_root, block_root);
+        self.fork_choice.write().add_block(&block, &block_root)?;
 
         // If the parent block was the parent_block, automatically update the canonical head.
         //
         // TODO: this is a first-in-best-dressed scenario that is not ideal; fork_choice should be
         // run instead.
         if self.head().beacon_block_root == parent_block_root {
-            self.update_canonical_head(block.clone(), block_root, state.clone(), state_root);
+            self.update_canonical_head(
+                block.clone(),
+                block_root.clone(),
+                state.clone(),
+                state_root,
+            );
             // Update the local state variable.
             *self.state.write() = state.clone();
         }
@@ -496,7 +496,7 @@ where
 
     /// Produce a new block at the present slot.
     ///
-    /// The produced block will not be inheriently valid, it must be signed by a block producer.
+    /// The produced block will not be inherently valid, it must be signed by a block producer.
     /// Block signing is out of the scope of this function and should be done by a separate program.
     pub fn produce_block(&self, randao_reveal: Signature) -> Option<(BeaconBlock, BeaconState)> {
         debug!("Producing block at slot {}...", self.state.read().slot);
@@ -549,6 +549,31 @@ where
 
         Some((block, state))
     }
+
+    // TODO: Left this as is, modify later
+    pub fn fork_choice(&self) -> Result<(), Error> {
+        let present_head = self.finalized_head().beacon_block_root;
+
+        let new_head = self.fork_choice.write().find_head(&present_head)?;
+
+        if new_head != present_head {
+            let block = self
+                .block_store
+                .get_deserialized(&new_head)?
+                .ok_or_else(|| Error::MissingBeaconBlock(new_head))?;
+            let block_root = block.canonical_root();
+
+            let state = self
+                .state_store
+                .get_deserialized(&block.state_root)?
+                .ok_or_else(|| Error::MissingBeaconState(block.state_root))?;
+            let state_root = state.canonical_root();
+
+            self.update_canonical_head(block, block_root, state, state_root);
+        }
+
+        Ok(())
+    }
 }
 
 impl From<DBError> for Error {
@@ -557,6 +582,12 @@ impl From<DBError> for Error {
     }
 }
 
+impl From<ForkChoiceError> for Error {
+    fn from(e: ForkChoiceError) -> Error {
+        Error::ForkChoiceError(e)
+    }
+}
+
 impl From<CommitteesError> for Error {
     fn from(e: CommitteesError) -> Error {
         Error::CommitteesError(e)

beacon_node/beacon_chain/src/block_graph.rs

@@ -1,45 +0,0 @@
-use parking_lot::{RwLock, RwLockReadGuard};
-use std::collections::HashSet;
-use types::Hash256;
-
-/// Maintains a view of the block DAG, also known as the "blockchain" (except, it tracks multiple
-/// chains eminating from a single root instead of just the head of some canonical chain).
-///
-/// The BlockGraph does not store the blocks, instead it tracks the block hashes of blocks at the
-/// tip of the DAG. It is out of the scope of the object to retrieve blocks.
-///
-/// Presently, the DAG root (genesis block) is not tracked.
-///
-/// The BlogGraph is thread-safe due to internal RwLocks.
-#[derive(Default)]
-pub struct BlockGraph {
-    pub leaves: RwLock<HashSet<Hash256>>,
-}
-
-impl BlockGraph {
-    /// Create a new block graph without any leaves.
-    pub fn new() -> Self {
-        Self {
-            leaves: RwLock::new(HashSet::new()),
-        }
-    }
-    /// Add a new leaf to the block hash graph. Returns `true` if the leaf was built upon another
-    /// leaf.
-    pub fn add_leaf(&self, parent: &Hash256, leaf: Hash256) -> bool {
-        let mut leaves = self.leaves.write();
-
-        if leaves.contains(parent) {
-            leaves.remove(parent);
-            leaves.insert(leaf);
-            true
-        } else {
-            leaves.insert(leaf);
-            false
-        }
-    }
-
-    /// Returns a read-guarded HashSet of all leaf blocks.
-    pub fn leaves(&self) -> RwLockReadGuard<HashSet<Hash256>> {
-        self.leaves.read()
-    }
-}

beacon_node/beacon_chain/src/lib.rs

@@ -1,9 +1,7 @@
 mod attestation_aggregator;
-mod attestation_targets;
 mod beacon_chain;
-mod block_graph;
 mod checkpoint;
-mod lmd_ghost;
 
 pub use self::beacon_chain::{BeaconChain, Error};
 pub use self::checkpoint::CheckPoint;
+pub use fork_choice::{ForkChoice, ForkChoiceAlgorithms, ForkChoiceError};

beacon_node/beacon_chain/src/lmd_ghost.rs

@@ -1,198 +0,0 @@
-use crate::BeaconChain;
-use db::{
-    stores::{BeaconBlockAtSlotError, BeaconBlockStore},
-    ClientDB, DBError,
-};
-use slot_clock::{SlotClock, TestingSlotClockError};
-use std::collections::HashSet;
-use std::sync::Arc;
-use types::{
-    readers::{BeaconBlockReader, BeaconStateReader},
-    validator_registry::get_active_validator_indices,
-    Hash256, Slot,
-};
-
-#[derive(Debug, PartialEq)]
-pub enum Error {
-    DBError(String),
-    MissingBeaconState(Hash256),
-    InvalidBeaconState(Hash256),
-    MissingBeaconBlock(Hash256),
-    InvalidBeaconBlock(Hash256),
-}
-
-impl<T, U> BeaconChain<T, U>
-where
-    T: ClientDB,
-    U: SlotClock,
-    Error: From<<U as SlotClock>::Error>,
-{
-    /// Run the fork-choice rule on the current chain, updating the canonical head, if required.
-    pub fn fork_choice(&self) -> Result<(), Error> {
-        let present_head = &self.finalized_head().beacon_block_root;
-
-        let new_head = self.slow_lmd_ghost(&self.finalized_head().beacon_block_root)?;
-
-        if new_head != *present_head {
-            let block = self
-                .block_store
-                .get_deserialized(&new_head)?
-                .ok_or_else(|| Error::MissingBeaconBlock(new_head))?;
-            let block_root = block.canonical_root();
-
-            let state = self
-                .state_store
-                .get_deserialized(&block.state_root)?
-                .ok_or_else(|| Error::MissingBeaconState(block.state_root))?;
-            let state_root = state.canonical_root();
-
-            self.update_canonical_head(block, block_root, state, state_root);
-        }
-
-        Ok(())
-    }
-
-    /// A very inefficient implementation of LMD ghost.
-    pub fn slow_lmd_ghost(&self, start_hash: &Hash256) -> Result<Hash256, Error> {
-        let start = self
-            .block_store
-            .get_reader(&start_hash)?
-            .ok_or_else(|| Error::MissingBeaconBlock(*start_hash))?;
-
-        let start_state_root = start.state_root();
-
-        let state = self
-            .state_store
-            .get_reader(&start_state_root)?
-            .ok_or_else(|| Error::MissingBeaconState(start_state_root))?
-            .into_beacon_state()
-            .ok_or_else(|| Error::InvalidBeaconState(start_state_root))?;
-
-        let active_validator_indices = get_active_validator_indices(
-            &state.validator_registry,
-            start.slot().epoch(self.spec.epoch_length),
-        );
-
-        let mut attestation_targets = Vec::with_capacity(active_validator_indices.len());
-        for i in active_validator_indices {
-            if let Some(target) = self.get_latest_attestation_target(i as u64) {
-                attestation_targets.push(target);
-            }
-        }
-
-        let mut head_hash = Hash256::zero();
-        let mut head_vote_count = 0;
-
-        loop {
-            let child_hashes_and_slots = get_child_hashes_and_slots(
-                &self.block_store,
-                &head_hash,
-                &self.block_graph.leaves(),
-            )?;
-
-            if child_hashes_and_slots.is_empty() {
-                break;
-            }
-
-            for (child_hash, child_slot) in child_hashes_and_slots {
-                let vote_count = get_vote_count(
-                    &self.block_store,
-                    &attestation_targets[..],
-                    &child_hash,
-                    child_slot,
-                )?;
-
-                if vote_count > head_vote_count {
-                    head_hash = child_hash;
-                    head_vote_count = vote_count;
-                }
-            }
-        }
-
-        Ok(head_hash)
-    }
-}
-
-/// Get the total number of votes for some given block root.
-///
-/// The vote count is incrememented each time an attestation target votes for a block root.
-fn get_vote_count<T: ClientDB>(
-    block_store: &Arc<BeaconBlockStore<T>>,
-    attestation_targets: &[Hash256],
-    block_root: &Hash256,
-    slot: Slot,
-) -> Result<u64, Error> {
-    let mut count = 0;
-    for target in attestation_targets {
-        let (root_at_slot, _) = block_store
-            .block_at_slot(&block_root, slot)?
-            .ok_or_else(|| Error::MissingBeaconBlock(*block_root))?;
-        if root_at_slot == *target {
-            count += 1;
-        }
-    }
-    Ok(count)
-}
-
-/// Starting from some `leaf_hashes`, recurse back down each branch until the `root_hash`, adding
-/// each `block_root` and `slot` to a HashSet.
-fn get_child_hashes_and_slots<T: ClientDB>(
-    block_store: &Arc<BeaconBlockStore<T>>,
-    root_hash: &Hash256,
-    leaf_hashes: &HashSet<Hash256>,
-) -> Result<HashSet<(Hash256, Slot)>, Error> {
-    let mut hash_set = HashSet::new();
-
-    for leaf_hash in leaf_hashes {
-        let mut current_hash = *leaf_hash;
-
-        loop {
-            if let Some(block_reader) = block_store.get_reader(&current_hash)? {
-                let parent_root = block_reader.parent_root();
-
-                let new_hash = hash_set.insert((current_hash, block_reader.slot()));
-
-                // If the hash just added was already in the set, break the loop.
-                //
-                // In such a case, the present branch has merged with a branch that is already in
-                // the set.
-                if !new_hash {
-                    break;
-                }
-
-                // The branch is exhausted if the parent of this block is the root_hash.
-                if parent_root == *root_hash {
-                    break;
-                }
-
-                current_hash = parent_root;
-            } else {
-                return Err(Error::MissingBeaconBlock(current_hash));
-            }
-        }
-    }
-
-    Ok(hash_set)
-}
-
-impl From<DBError> for Error {
-    fn from(e: DBError) -> Error {
-        Error::DBError(e.message)
-    }
-}
-
-impl From<BeaconBlockAtSlotError> for Error {
-    fn from(e: BeaconBlockAtSlotError) -> Error {
-        match e {
-            BeaconBlockAtSlotError::UnknownBeaconBlock(h) => Error::MissingBeaconBlock(h),
-            BeaconBlockAtSlotError::InvalidBeaconBlock(h) => Error::InvalidBeaconBlock(h),
-            BeaconBlockAtSlotError::DBError(msg) => Error::DBError(msg),
-        }
-    }
-}
-
-impl From<TestingSlotClockError> for Error {
-    fn from(_: TestingSlotClockError) -> Error {
-        unreachable!(); // Testing clock never throws an error.
-    }
-}

beacon_node/beacon_chain/test_harness/Cargo.toml

@@ -21,6 +21,7 @@ db = { path = "../../db" }
 parking_lot = "0.7"
 failure = "0.1"
 failure_derive = "0.1"
+fork_choice = { path = "../../../eth2/fork_choice" }
 hashing = { path = "../../../eth2/utils/hashing" }
 log = "0.4"
 env_logger = "0.6.0"

beacon_node/beacon_chain/test_harness/src/beacon_chain_harness.rs

@@ -6,6 +6,7 @@ use db::{
     stores::{BeaconBlockStore, BeaconStateStore},
     MemoryDB,
 };
+use fork_choice::{optimised_lmd_ghost::OptimisedLMDGhost, slow_lmd_ghost::SlowLMDGhost}; // import all the algorithms
 use log::debug;
 use rayon::prelude::*;
 use slot_clock::TestingSlotClock;
@@ -21,13 +22,13 @@ use types::{
 
 /// The beacon chain harness simulates a single beacon node with `validator_count` validators connected
 /// to it. Each validator is provided a borrow to the beacon chain, where it may read
-/// information and submit blocks/attesations for processing.
+/// information and submit blocks/attestations for processing.
 ///
 /// This test harness is useful for testing validator and internal state transition logic. It
 /// is not useful for testing that multiple beacon nodes can reach consensus.
 pub struct BeaconChainHarness {
     pub db: Arc<MemoryDB>,
-    pub beacon_chain: Arc<BeaconChain<MemoryDB, TestingSlotClock>>,
+    pub beacon_chain: Arc<BeaconChain<MemoryDB, TestingSlotClock, OptimisedLMDGhost<MemoryDB>>>,
     pub block_store: Arc<BeaconBlockStore<MemoryDB>>,
     pub state_store: Arc<BeaconStateStore<MemoryDB>>,
     pub validators: Vec<ValidatorHarness>,
@@ -43,9 +44,9 @@ impl BeaconChainHarness {
         let db = Arc::new(MemoryDB::open());
         let block_store = Arc::new(BeaconBlockStore::new(db.clone()));
         let state_store = Arc::new(BeaconStateStore::new(db.clone()));
-
         let genesis_time = 1_549_935_547; // 12th Feb 2018 (arbitrary value in the past).
         let slot_clock = TestingSlotClock::new(spec.genesis_slot.as_u64());
+        let fork_choice = OptimisedLMDGhost::new(block_store.clone(), state_store.clone());
         let latest_eth1_data = Eth1Data {
             deposit_root: Hash256::zero(),
             block_hash: Hash256::zero(),
@@ -90,6 +91,7 @@ impl BeaconChainHarness {
                 latest_eth1_data,
                 initial_validator_deposits,
                 spec.clone(),
+                fork_choice,
             )
             .unwrap(),
         );

beacon_node/beacon_chain/test_harness/src/validator_harness/direct_beacon_node.rs

@@ -8,6 +8,7 @@ use block_producer::{
     PublishOutcome as BlockPublishOutcome,
 };
 use db::ClientDB;
+use fork_choice::ForkChoice;
 use parking_lot::RwLock;
 use slot_clock::SlotClock;
 use std::sync::Arc;
@@ -22,14 +23,14 @@ use types::{AttestationData, BeaconBlock, FreeAttestation, Signature, Slot};
 /// `BeaconBlock`s and `FreeAttestation`s are not actually published to the `BeaconChain`, instead
 /// they are stored inside this struct. This is to allow one to benchmark the submission of the
 /// block/attestation directly, or modify it before submission.
-pub struct DirectBeaconNode<T: ClientDB, U: SlotClock> {
+pub struct DirectBeaconNode<T: ClientDB, U: SlotClock, F: ForkChoice> {
-    beacon_chain: Arc<BeaconChain<T, U>>,
+    beacon_chain: Arc<BeaconChain<T, U, F>>,
     published_blocks: RwLock<Vec<BeaconBlock>>,
     published_attestations: RwLock<Vec<FreeAttestation>>,
 }
 
-impl<T: ClientDB, U: SlotClock> DirectBeaconNode<T, U> {
+impl<T: ClientDB, U: SlotClock, F: ForkChoice> DirectBeaconNode<T, U, F> {
-    pub fn new(beacon_chain: Arc<BeaconChain<T, U>>) -> Self {
+    pub fn new(beacon_chain: Arc<BeaconChain<T, U, F>>) -> Self {
         Self {
             beacon_chain,
             published_blocks: RwLock::new(vec![]),
@@ -48,7 +49,7 @@ impl<T: ClientDB, U: SlotClock> DirectBeaconNode<T, U> {
     }
 }
 
-impl<T: ClientDB, U: SlotClock> AttesterBeaconNode for DirectBeaconNode<T, U> {
+impl<T: ClientDB, U: SlotClock, F: ForkChoice> AttesterBeaconNode for DirectBeaconNode<T, U, F> {
     fn produce_attestation_data(
         &self,
         _slot: Slot,
@@ -69,7 +70,7 @@ impl<T: ClientDB, U: SlotClock> AttesterBeaconNode for DirectBeaconNode<T, U> {
     }
 }
 
-impl<T: ClientDB, U: SlotClock> BeaconBlockNode for DirectBeaconNode<T, U> {
+impl<T: ClientDB, U: SlotClock, F: ForkChoice> BeaconBlockNode for DirectBeaconNode<T, U, F> {
     /// Requests a new `BeaconBlock from the `BeaconChain`.
     fn produce_beacon_block(
         &self,

beacon_node/beacon_chain/test_harness/src/validator_harness/direct_duties.rs

@@ -6,19 +6,20 @@ use block_producer::{
     DutiesReader as ProducerDutiesReader, DutiesReaderError as ProducerDutiesReaderError,
 };
 use db::ClientDB;
+use fork_choice::ForkChoice;
 use slot_clock::SlotClock;
 use std::sync::Arc;
 use types::{PublicKey, Slot};
 
 /// Connects directly to a borrowed `BeaconChain` and reads attester/proposer duties directly from
 /// it.
-pub struct DirectDuties<T: ClientDB, U: SlotClock> {
+pub struct DirectDuties<T: ClientDB, U: SlotClock, F: ForkChoice> {
-    beacon_chain: Arc<BeaconChain<T, U>>,
+    beacon_chain: Arc<BeaconChain<T, U, F>>,
     pubkey: PublicKey,
 }
 
-impl<T: ClientDB, U: SlotClock> DirectDuties<T, U> {
+impl<T: ClientDB, U: SlotClock, F: ForkChoice> DirectDuties<T, U, F> {
-    pub fn new(pubkey: PublicKey, beacon_chain: Arc<BeaconChain<T, U>>) -> Self {
+    pub fn new(pubkey: PublicKey, beacon_chain: Arc<BeaconChain<T, U, F>>) -> Self {
         Self {
             beacon_chain,
             pubkey,
@@ -26,7 +27,7 @@ impl<T: ClientDB, U: SlotClock> DirectDuties<T, U> {
     }
 }
 
-impl<T: ClientDB, U: SlotClock> ProducerDutiesReader for DirectDuties<T, U> {
+impl<T: ClientDB, U: SlotClock, F: ForkChoice> ProducerDutiesReader for DirectDuties<T, U, F> {
     fn is_block_production_slot(&self, slot: Slot) -> Result<bool, ProducerDutiesReaderError> {
         let validator_index = self
             .beacon_chain
@@ -41,7 +42,7 @@ impl<T: ClientDB, U: SlotClock> ProducerDutiesReader for DirectDuties<T, U> {
     }
 }
 
-impl<T: ClientDB, U: SlotClock> AttesterDutiesReader for DirectDuties<T, U> {
+impl<T: ClientDB, U: SlotClock, F: ForkChoice> AttesterDutiesReader for DirectDuties<T, U, F> {
     fn validator_index(&self) -> Option<u64> {
         match self.beacon_chain.validator_index(&self.pubkey) {
             Some(index) => Some(index as u64),

beacon_node/beacon_chain/test_harness/src/validator_harness/mod.rs

@@ -10,6 +10,7 @@ use block_producer::{BlockProducer, Error as BlockPollError};
 use db::MemoryDB;
 use direct_beacon_node::DirectBeaconNode;
 use direct_duties::DirectDuties;
+use fork_choice::{optimised_lmd_ghost::OptimisedLMDGhost, slow_lmd_ghost::SlowLMDGhost};
 use local_signer::LocalSigner;
 use slot_clock::TestingSlotClock;
 use std::sync::Arc;
@@ -35,20 +36,20 @@ pub enum AttestationProduceError {
 pub struct ValidatorHarness {
     pub block_producer: BlockProducer<
         TestingSlotClock,
-        DirectBeaconNode<MemoryDB, TestingSlotClock>,
+        DirectBeaconNode<MemoryDB, TestingSlotClock, OptimisedLMDGhost<MemoryDB>>,
-        DirectDuties<MemoryDB, TestingSlotClock>,
+        DirectDuties<MemoryDB, TestingSlotClock, OptimisedLMDGhost<MemoryDB>>,
         LocalSigner,
     >,
     pub attester: Attester<
         TestingSlotClock,
-        DirectBeaconNode<MemoryDB, TestingSlotClock>,
+        DirectBeaconNode<MemoryDB, TestingSlotClock, OptimisedLMDGhost<MemoryDB>>,
-        DirectDuties<MemoryDB, TestingSlotClock>,
+        DirectDuties<MemoryDB, TestingSlotClock, OptimisedLMDGhost<MemoryDB>>,
         LocalSigner,
     >,
     pub spec: Arc<ChainSpec>,
-    pub epoch_map: Arc<DirectDuties<MemoryDB, TestingSlotClock>>,
+    pub epoch_map: Arc<DirectDuties<MemoryDB, TestingSlotClock, OptimisedLMDGhost<MemoryDB>>>,
     pub keypair: Keypair,
-    pub beacon_node: Arc<DirectBeaconNode<MemoryDB, TestingSlotClock>>,
+    pub beacon_node: Arc<DirectBeaconNode<MemoryDB, TestingSlotClock, OptimisedLMDGhost<MemoryDB>>>,
     pub slot_clock: Arc<TestingSlotClock>,
     pub signer: Arc<LocalSigner>,
 }
@@ -60,7 +61,7 @@ impl ValidatorHarness {
     /// A `BlockProducer` and `Attester` is created..
     pub fn new(
         keypair: Keypair,
-        beacon_chain: Arc<BeaconChain<MemoryDB, TestingSlotClock>>,
+        beacon_chain: Arc<BeaconChain<MemoryDB, TestingSlotClock, OptimisedLMDGhost<MemoryDB>>>,
         spec: Arc<ChainSpec>,
     ) -> Self {
         let slot_clock = Arc::new(TestingSlotClock::new(spec.genesis_slot.as_u64()));

beacon_node/src/main.rs

@@ -14,6 +14,7 @@ use db::{
     stores::{BeaconBlockStore, BeaconStateStore},
     MemoryDB,
 };
+use fork_choice::optimised_lmd_ghost::OptimisedLMDGhost;
 use slog::{error, info, o, Drain};
 use slot_clock::SystemTimeSlotClock;
 use std::sync::Arc;
@@ -79,6 +80,8 @@ fn main() {
     let genesis_time = 1_549_935_547; // 12th Feb 2018 (arbitrary value in the past).
     let slot_clock = SystemTimeSlotClock::new(genesis_time, spec.slot_duration)
         .expect("Unable to load SystemTimeSlotClock");
+    // Choose the fork choice
+    let fork_choice = OptimisedLMDGhost::new(block_store.clone(), state_store.clone());
 
     /*
      * Generate some random data to start a chain with.
@@ -120,6 +123,7 @@ fn main() {
         latest_eth1_data,
         initial_validator_deposits,
         spec,
+        fork_choice,
     );
 
     let _server = start_server(log.clone());

eth2/fork_choice/Cargo.toml

@@ -1,10 +1,12 @@
 [package]
-name = "naive_fork_choice"
+name = "fork_choice"
 version = "0.1.0"
-authors = ["Paul Hauner <paul@paulhauner.com>"]
+authors = ["Age Manning <Age@AgeManning.com>"]
 edition = "2018"
 
 [dependencies]
 db = { path = "../../beacon_node/db" }
 ssz = { path = "../utils/ssz" }
 types = { path = "../types" }
+fast-math = "0.1.1"
+byteorder = "1.3.1"

eth2/fork_choice/src/lib.rs

@@ -0,0 +1,118 @@
+// Copyright 2019 Sigma Prime Pty Ltd.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+//! 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
+//! purposes). The other three 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
+//! 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.
+//!
+//! [`slow_lmd_ghost`]: struct.SlowLmdGhost.html
+//! [`optimised_lmd_ghost`]: struct.OptimisedLmdGhost.html
+//! [`protolambda_lmd_ghost`]: struct.ProtolambdaLmdGhost.html
+
+extern crate db;
+extern crate ssz;
+extern crate types;
+
+pub mod longest_chain;
+pub mod optimised_lmd_ghost;
+pub mod protolambda_lmd_ghost;
+pub mod slow_lmd_ghost;
+
+use db::stores::BeaconBlockAtSlotError;
+use db::DBError;
+use types::{BeaconBlock, Hash256};
+
+/// 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.
+/// The main fork choice algorithm is specified in `find_head
+pub trait ForkChoice: Send + Sync {
+    /// Called when a block has been added. Allows generic block-level data structures to be
+    /// built for a given fork-choice.
+    fn add_block(
+        &mut self,
+        block: &BeaconBlock,
+        block_hash: &Hash256,
+    ) -> Result<(), ForkChoiceError>;
+    /// Called when an attestation has been added. Allows generic attestation-level data structures to be built for a given fork choice.
+    // This can be generalised to a full attestation if required later.
+    fn add_attestation(
+        &mut self,
+        validator_index: u64,
+        target_block_hash: &Hash256,
+    ) -> Result<(), ForkChoiceError>;
+    /// The fork-choice algorithm to find the current canonical head of the chain.
+    // TODO: Remove the justified_start_block parameter and make it internal
+    fn find_head(&mut self, justified_start_block: &Hash256) -> Result<Hash256, ForkChoiceError>;
+}
+
+/// Possible fork choice errors that can occur.
+#[derive(Debug, PartialEq)]
+pub enum ForkChoiceError {
+    MissingBeaconBlock(Hash256),
+    MissingBeaconState(Hash256),
+    IncorrectBeaconState(Hash256),
+    CannotFindBestChild,
+    ChildrenNotFound,
+    StorageError(String),
+}
+
+impl From<DBError> for ForkChoiceError {
+    fn from(e: DBError) -> ForkChoiceError {
+        ForkChoiceError::StorageError(e.message)
+    }
+}
+
+impl From<BeaconBlockAtSlotError> for ForkChoiceError {
+    fn from(e: BeaconBlockAtSlotError) -> ForkChoiceError {
+        match e {
+            BeaconBlockAtSlotError::UnknownBeaconBlock(hash) => {
+                ForkChoiceError::MissingBeaconBlock(hash)
+            }
+            BeaconBlockAtSlotError::InvalidBeaconBlock(hash) => {
+                ForkChoiceError::MissingBeaconBlock(hash)
+            }
+            BeaconBlockAtSlotError::DBError(string) => ForkChoiceError::StorageError(string),
+        }
+    }
+}
+
+/// Fork choice options that are currently implemented.
+pub enum ForkChoiceAlgorithms {
+    /// Chooses the longest chain becomes the head. Not for production.
+    LongestChain,
+    /// A simple and highly inefficient implementation of LMD ghost.
+    SlowLMDGhost,
+    /// An optimised version of LMD-GHOST by Vitalik.
+    OptimisedLMDGhost,
+    /// An optimised version of LMD-GHOST by Protolambda.
+    ProtoLMDGhost,
+}

eth2/fork_choice/src/longest_chain.rs

@@ -1,7 +1,3 @@
-extern crate db;
-extern crate ssz;
-extern crate types;
-
 use db::stores::BeaconBlockStore;
 use db::{ClientDB, DBError};
 use ssz::{Decodable, DecodeError};
@@ -14,7 +10,7 @@ pub enum ForkChoiceError {
     DBError(String),
 }
 
-pub fn naive_fork_choice<T>(
+pub fn longest_chain<T>(
     head_block_hashes: &[Hash256],
     block_store: &Arc<BeaconBlockStore<T>>,
 ) -> Result<Option<usize>, ForkChoiceError>

eth2/fork_choice/src/optimised_lmd_ghost.rs

@@ -0,0 +1,443 @@
+// Copyright 2019 Sigma Prime Pty Ltd.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+extern crate byteorder;
+extern crate fast_math;
+use crate::{ForkChoice, ForkChoiceError};
+use byteorder::{BigEndian, ByteOrder};
+use db::{
+    stores::{BeaconBlockStore, BeaconStateStore},
+    ClientDB,
+};
+use fast_math::log2_raw;
+use std::collections::HashMap;
+use std::sync::Arc;
+use types::{
+    readers::BeaconBlockReader,
+    slot_epoch_height::{Height, Slot},
+    validator_registry::get_active_validator_indices,
+    BeaconBlock, Hash256,
+};
+
+//TODO: Pruning - Children
+//TODO: Handle Syncing
+
+//TODO: Sort out global constants
+const GENESIS_SLOT: u64 = 0;
+const FORK_CHOICE_BALANCE_INCREMENT: u64 = 1e9 as u64;
+const MAX_DEPOSIT_AMOUNT: u64 = 32e9 as u64;
+const EPOCH_LENGTH: u64 = 64;
+
+/// The optimised LMD-GHOST fork choice rule.
+/// NOTE: This uses u32 to represent difference between block heights. Thus this is only
+/// applicable for block height differences in the range of a u32.
+/// This can potentially be parallelized in some parts.
+// we use fast log2, a log2 lookup table is implemented in Vitaliks code, potentially do
+// the comparison. Log2_raw takes 2ns according to the documentation.
+#[inline]
+fn log2_int(x: u32) -> u32 {
+    log2_raw(x as f32) as u32
+}
+
+fn power_of_2_below(x: u32) -> u32 {
+    2u32.pow(log2_int(x))
+}
+
+/// Stores the necessary data structures to run the optimised lmd ghost algorithm.
+pub struct OptimisedLMDGhost<T: ClientDB + Sized> {
+    /// A cache of known ancestors at given heights for a specific block.
+    //TODO: Consider FnvHashMap
+    cache: HashMap<CacheKey<u32>, Hash256>,
+    /// Log lookup table for blocks to their ancestors.
+    //TODO: Verify we only want/need a size 16 log lookup
+    ancestors: Vec<HashMap<Hash256, Hash256>>,
+    /// Stores the children for any given parent.
+    children: HashMap<Hash256, Vec<Hash256>>,
+    /// The latest attestation targets as a map of validator index to block hash.
+    //TODO: Could this be a fixed size vec
+    latest_attestation_targets: HashMap<u64, Hash256>,
+    /// Block storage access.
+    block_store: Arc<BeaconBlockStore<T>>,
+    /// State storage access.
+    state_store: Arc<BeaconStateStore<T>>,
+    max_known_height: Height,
+}
+
+impl<T> OptimisedLMDGhost<T>
+where
+    T: ClientDB + Sized,
+{
+    pub fn new(
+        block_store: Arc<BeaconBlockStore<T>>,
+        state_store: Arc<BeaconStateStore<T>>,
+    ) -> Self {
+        OptimisedLMDGhost {
+            cache: HashMap::new(),
+            ancestors: vec![HashMap::new(); 16],
+            latest_attestation_targets: HashMap::new(),
+            children: HashMap::new(),
+            max_known_height: Height::new(0),
+            block_store,
+            state_store,
+        }
+    }
+
+    /// Finds the latest votes weighted by validator balance. Returns a hashmap of block_hash to
+    /// weighted votes.
+    pub fn get_latest_votes(
+        &self,
+        state_root: &Hash256,
+        block_slot: Slot,
+    ) -> Result<HashMap<Hash256, u64>, ForkChoiceError> {
+        // get latest votes
+        // Note: Votes are weighted by min(balance, MAX_DEPOSIT_AMOUNT) //
+        // FORK_CHOICE_BALANCE_INCREMENT
+        // build a hashmap of block_hash to weighted votes
+        let mut latest_votes: HashMap<Hash256, u64> = HashMap::new();
+        // gets the current weighted votes
+        let current_state = self
+            .state_store
+            .get_deserialized(&state_root)?
+            .ok_or_else(|| ForkChoiceError::MissingBeaconState(*state_root))?;
+
+        let active_validator_indices = get_active_validator_indices(
+            &current_state.validator_registry,
+            block_slot.epoch(EPOCH_LENGTH),
+        );
+
+        for index in active_validator_indices {
+            let balance =
+                std::cmp::min(current_state.validator_balances[index], MAX_DEPOSIT_AMOUNT)
+                    / 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;
+                }
+            }
+        }
+
+        Ok(latest_votes)
+    }
+
+    /// Gets the ancestor at a given height `at_height` of a block specified by `block_hash`.
+    fn get_ancestor(&mut self, block_hash: Hash256, at_height: Height) -> Option<Hash256> {
+        // return None if we can't get the block from the db.
+        let block_height = {
+            let block_slot = self
+                .block_store
+                .get_deserialized(&block_hash)
+                .ok()?
+                .expect("Should have returned already if None")
+                .slot;
+
+            block_slot.height(Slot::from(GENESIS_SLOT))
+        };
+
+        // verify we haven't exceeded the block height
+        if at_height >= block_height {
+            if at_height > block_height {
+                return None;
+            } else {
+                return Some(block_hash);
+            }
+        }
+        // check if the result is stored in our cache
+        let cache_key = CacheKey::new(&block_hash, at_height.as_u32());
+        if let Some(ancestor) = self.cache.get(&cache_key) {
+            return Some(*ancestor);
+        }
+
+        // not in the cache recursively search for ancestors using a log-lookup
+
+        if let Some(ancestor) = {
+            let ancestor_lookup = self.ancestors
+                [log2_int((block_height - at_height - 1u64).as_u32()) as usize]
+                .get(&block_hash)
+                //TODO: Panic if we can't lookup and fork choice fails
+                .expect("All blocks should be added to the ancestor log lookup table");
+            self.get_ancestor(*ancestor_lookup, at_height)
+        } {
+            // add the result to the cache
+            self.cache.insert(cache_key, ancestor);
+            return Some(ancestor);
+        }
+
+        None
+    }
+
+    // looks for an obvious block winner given the latest votes for a specific height
+    fn get_clear_winner(
+        &mut self,
+        latest_votes: &HashMap<Hash256, u64>,
+        block_height: Height,
+    ) -> Option<Hash256> {
+        // map of vote counts for every hash at this height
+        let mut current_votes: HashMap<Hash256, u64> = HashMap::new();
+        let mut total_vote_count = 0;
+
+        // loop through the latest votes and count all votes
+        // these have already been weighted by balance
+        for (hash, votes) in latest_votes.iter() {
+            if let Some(ancestor) = self.get_ancestor(*hash, block_height) {
+                let current_vote_value = current_votes.get(&ancestor).unwrap_or_else(|| &0);
+                current_votes.insert(ancestor, current_vote_value + *votes);
+                total_vote_count += votes;
+            }
+        }
+        // Check if there is a clear block winner at this height. If so return it.
+        for (hash, votes) in current_votes.iter() {
+            if *votes >= total_vote_count / 2 {
+                // we have a clear winner, return it
+                return Some(*hash);
+            }
+        }
+        // didn't find a clear winner
+        None
+    }
+
+    // Finds the best child, splitting children into a binary tree, based on their hashes
+    fn choose_best_child(&self, votes: &HashMap<Hash256, u64>) -> Option<Hash256> {
+        let mut bitmask = 0;
+        for bit in (0..=255).rev() {
+            let mut zero_votes = 0;
+            let mut one_votes = 0;
+            let mut single_candidate = None;
+
+            for (candidate, votes) in votes.iter() {
+                let candidate_uint = BigEndian::read_u32(candidate);
+                if candidate_uint >> (bit + 1) != bitmask {
+                    continue;
+                }
+                if (candidate_uint >> bit) % 2 == 0 {
+                    zero_votes += votes;
+                } else {
+                    one_votes += votes;
+                }
+
+                if single_candidate.is_none() {
+                    single_candidate = Some(candidate);
+                } else {
+                    single_candidate = None;
+                }
+            }
+            bitmask = (bitmask * 2) + {
+                if one_votes > zero_votes {
+                    1
+                } else {
+                    0
+                }
+            };
+            if let Some(candidate) = single_candidate {
+                return Some(*candidate);
+            }
+            //TODO Remove this during benchmark after testing
+            assert!(bit >= 1);
+        }
+        // should never reach here
+        None
+    }
+}
+
+impl<T: ClientDB + Sized> ForkChoice for OptimisedLMDGhost<T> {
+    fn add_block(
+        &mut self,
+        block: &BeaconBlock,
+        block_hash: &Hash256,
+    ) -> Result<(), ForkChoiceError> {
+        // get the height of the parent
+        let parent_height = self
+            .block_store
+            .get_deserialized(&block.parent_root)?
+            .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(block.parent_root))?
+            .slot()
+            .height(Slot::from(GENESIS_SLOT));
+
+        let parent_hash = &block.parent_root;
+
+        // add the new block to the children of parent
+        (*self
+            .children
+            .entry(block.parent_root)
+            .or_insert_with(|| vec![]))
+        .push(block_hash.clone());
+
+        // build the ancestor data structure
+        for index in 0..16 {
+            if parent_height % (1 << index) == 0 {
+                self.ancestors[index].insert(*block_hash, *parent_hash);
+            } else {
+                // TODO: This is unsafe. Will panic if parent_hash doesn't exist. Using it for debugging
+                let parent_ancestor = self.ancestors[index][parent_hash];
+                self.ancestors[index].insert(*block_hash, parent_ancestor);
+            }
+        }
+        // update the max height
+        self.max_known_height = std::cmp::max(self.max_known_height, parent_height + 1);
+        Ok(())
+    }
+
+    fn add_attestation(
+        &mut self,
+        validator_index: u64,
+        target_block_root: &Hash256,
+    ) -> Result<(), ForkChoiceError> {
+        // simply add the attestation to the latest_attestation_target if the block_height is
+        // larger
+        let attestation_target = self
+            .latest_attestation_targets
+            .entry(validator_index)
+            .or_insert_with(|| *target_block_root);
+        // if we already have a value
+        if attestation_target != target_block_root {
+            // get the height of the target block
+            let block_height = self
+                .block_store
+                .get_deserialized(&target_block_root)?
+                .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(*target_block_root))?
+                .slot()
+                .height(Slot::from(GENESIS_SLOT));
+
+            // get the height of the past target block
+            let past_block_height = self
+                .block_store
+                .get_deserialized(&attestation_target)?
+                .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(*attestation_target))?
+                .slot()
+                .height(Slot::from(GENESIS_SLOT));
+            // update the attestation only if the new target is higher
+            if past_block_height < block_height {
+                *attestation_target = *target_block_root;
+            }
+        }
+        Ok(())
+    }
+
+    /// Perform lmd_ghost on the current chain to find the head.
+    fn find_head(&mut self, justified_block_start: &Hash256) -> Result<Hash256, ForkChoiceError> {
+        let block = self
+            .block_store
+            .get_deserialized(&justified_block_start)?
+            .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(*justified_block_start))?;
+
+        let block_slot = block.slot();
+        let block_height = block_slot.height(Slot::from(GENESIS_SLOT));
+        let state_root = block.state_root();
+
+        let mut current_head = *justified_block_start;
+
+        let mut latest_votes = self.get_latest_votes(&state_root, block_slot)?;
+
+        // remove any votes that don't relate to our current head.
+        latest_votes.retain(|hash, _| self.get_ancestor(*hash, block_height) == Some(current_head));
+
+        // begin searching for the head
+        loop {
+            // if there are no children, we are done, return the current_head
+            let children = match self.children.get(&current_head) {
+                Some(children) => children.clone(),
+                None => return Ok(current_head),
+            };
+
+            // logarithmic lookup blocks to see if there are obvious winners, if so,
+            // progress to the next iteration.
+            let mut step =
+                power_of_2_below(self.max_known_height.saturating_sub(block_height).as_u32()) / 2;
+            while step > 0 {
+                if let Some(clear_winner) = self.get_clear_winner(
+                    &latest_votes,
+                    block_height - (block_height % u64::from(step)) + u64::from(step),
+                ) {
+                    current_head = clear_winner;
+                    break;
+                }
+                step /= 2;
+            }
+            if step > 0 {
+            }
+            // if our skip lookup failed and we only have one child, progress to that child
+            else if children.len() == 1 {
+                current_head = children[0];
+            }
+            // we need to find the best child path to progress down.
+            else {
+                let mut child_votes = HashMap::new();
+                for (voted_hash, vote) in latest_votes.iter() {
+                    // if the latest votes correspond to a child
+                    if let Some(child) = self.get_ancestor(*voted_hash, block_height + 1) {
+                        // add up the votes for each child
+                        *child_votes.entry(child).or_insert_with(|| 0) += vote;
+                    }
+                }
+                // given the votes on the children, find the best child
+                current_head = self
+                    .choose_best_child(&child_votes)
+                    .ok_or(ForkChoiceError::CannotFindBestChild)?;
+            }
+
+            // No head was found, re-iterate
+
+            // update the block height for the next iteration
+            let block_height = self
+                .block_store
+                .get_deserialized(&current_head)?
+                .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(*justified_block_start))?
+                .slot()
+                .height(Slot::from(GENESIS_SLOT));
+
+            // prune the latest votes for votes that are not part of current chosen chain
+            // more specifically, only keep votes that have head as an ancestor
+            latest_votes
+                .retain(|hash, _| self.get_ancestor(*hash, block_height) == Some(current_head));
+        }
+    }
+}
+
+/// Type for storing blocks in a memory cache. Key is comprised of block-hash plus the height.
+#[derive(PartialEq, Eq, Hash)]
+pub struct CacheKey<T> {
+    block_hash: Hash256,
+    block_height: T,
+}
+
+impl<T> CacheKey<T> {
+    pub fn new(block_hash: &Hash256, block_height: T) -> Self {
+        CacheKey {
+            block_hash: *block_hash,
+            block_height,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    pub fn test_power_of_2_below() {
+        println!("{:?}", std::f32::MAX);
+        assert_eq!(power_of_2_below(4), 4);
+        assert_eq!(power_of_2_below(5), 4);
+        assert_eq!(power_of_2_below(7), 4);
+        assert_eq!(power_of_2_below(24), 16);
+        assert_eq!(power_of_2_below(32), 32);
+        assert_eq!(power_of_2_below(33), 32);
+        assert_eq!(power_of_2_below(63), 32);
+    }
+}

eth2/fork_choice/src/slow_lmd_ghost.rs

@@ -0,0 +1,223 @@
+// Copyright 2019 Sigma Prime Pty Ltd.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+extern crate db;
+
+use crate::{ForkChoice, ForkChoiceError};
+use db::{
+    stores::{BeaconBlockStore, BeaconStateStore},
+    ClientDB,
+};
+use std::collections::HashMap;
+use std::sync::Arc;
+use types::{
+    readers::{BeaconBlockReader, BeaconStateReader},
+    slot_epoch_height::Slot,
+    validator_registry::get_active_validator_indices,
+    BeaconBlock, Hash256,
+};
+
+//TODO: Pruning and syncing
+
+//TODO: Sort out global constants
+const GENESIS_SLOT: u64 = 0;
+const FORK_CHOICE_BALANCE_INCREMENT: u64 = 1e9 as u64;
+const MAX_DEPOSIT_AMOUNT: u64 = 32e9 as u64;
+const EPOCH_LENGTH: u64 = 64;
+
+pub struct SlowLMDGhost<T: ClientDB + Sized> {
+    /// The latest attestation targets as a map of validator index to block hash.
+    //TODO: Could this be a fixed size vec
+    latest_attestation_targets: HashMap<u64, Hash256>,
+    /// Stores the children for any given parent.
+    children: HashMap<Hash256, Vec<Hash256>>,
+    /// Block storage access.
+    block_store: Arc<BeaconBlockStore<T>>,
+    /// State storage access.
+    state_store: Arc<BeaconStateStore<T>>,
+}
+
+impl<T> SlowLMDGhost<T>
+where
+    T: ClientDB + Sized,
+{
+    pub fn new(block_store: BeaconBlockStore<T>, state_store: BeaconStateStore<T>) -> Self {
+        SlowLMDGhost {
+            latest_attestation_targets: HashMap::new(),
+            children: HashMap::new(),
+            block_store: Arc::new(block_store),
+            state_store: Arc::new(state_store),
+        }
+    }
+
+    /// Finds the latest votes weighted by validator balance. Returns a hashmap of block_hash to
+    /// weighted votes.
+    pub fn get_latest_votes(
+        &self,
+        state_root: &Hash256,
+        block_slot: Slot,
+    ) -> Result<HashMap<Hash256, u64>, ForkChoiceError> {
+        // get latest votes
+        // Note: Votes are weighted by min(balance, MAX_DEPOSIT_AMOUNT) //
+        // FORK_CHOICE_BALANCE_INCREMENT
+        // build a hashmap of block_hash to weighted votes
+        let mut latest_votes: HashMap<Hash256, u64> = HashMap::new();
+        // gets the current weighted votes
+        let current_state = self
+            .state_store
+            .get_deserialized(&state_root)?
+            .ok_or_else(|| ForkChoiceError::MissingBeaconState(*state_root))?;
+
+        let active_validator_indices = get_active_validator_indices(
+            &current_state.validator_registry,
+            block_slot.epoch(EPOCH_LENGTH),
+        );
+
+        for index in active_validator_indices {
+            let balance =
+                std::cmp::min(current_state.validator_balances[index], MAX_DEPOSIT_AMOUNT)
+                    / 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;
+                }
+            }
+        }
+
+        Ok(latest_votes)
+    }
+
+    /// Get the total number of votes for some given block root.
+    ///
+    /// The vote count is incremented each time an attestation target votes for a block root.
+    fn get_vote_count(
+        &self,
+        latest_votes: &HashMap<Hash256, u64>,
+        block_root: &Hash256,
+    ) -> Result<u64, ForkChoiceError> {
+        let mut count = 0;
+        let block_slot = self
+            .block_store
+            .get_deserialized(&block_root)?
+            .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(*block_root))?
+            .slot();
+
+        for (target_hash, votes) in latest_votes.iter() {
+            let (root_at_slot, _) = self
+                .block_store
+                .block_at_slot(&block_root, block_slot)?
+                .ok_or(ForkChoiceError::MissingBeaconBlock(*block_root))?;
+            if root_at_slot == *target_hash {
+                count += votes;
+            }
+        }
+        Ok(count)
+    }
+}
+
+impl<T: ClientDB + Sized> ForkChoice for SlowLMDGhost<T> {
+    /// Process when a block is added
+    fn add_block(
+        &mut self,
+        block: &BeaconBlock,
+        block_hash: &Hash256,
+    ) -> Result<(), ForkChoiceError> {
+        // build the children hashmap
+        // add the new block to the children of parent
+        (*self
+            .children
+            .entry(block.parent_root)
+            .or_insert_with(|| vec![]))
+        .push(block_hash.clone());
+
+        // complete
+        Ok(())
+    }
+
+    fn add_attestation(
+        &mut self,
+        validator_index: u64,
+        target_block_root: &Hash256,
+    ) -> Result<(), ForkChoiceError> {
+        // simply add the attestation to the latest_attestation_target if the block_height is
+        // larger
+        let attestation_target = self
+            .latest_attestation_targets
+            .entry(validator_index)
+            .or_insert_with(|| *target_block_root);
+        // if we already have a value
+        if attestation_target != target_block_root {
+            // get the height of the target block
+            let block_height = self
+                .block_store
+                .get_deserialized(&target_block_root)?
+                .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(*target_block_root))?
+                .slot()
+                .height(Slot::from(GENESIS_SLOT));
+
+            // get the height of the past target block
+            let past_block_height = self
+                .block_store
+                .get_deserialized(&attestation_target)?
+                .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(*attestation_target))?
+                .slot()
+                .height(Slot::from(GENESIS_SLOT));
+            // update the attestation only if the new target is higher
+            if past_block_height < block_height {
+                *attestation_target = *target_block_root;
+            }
+        }
+        Ok(())
+    }
+
+    /// A very inefficient implementation of LMD ghost.
+    fn find_head(&mut self, justified_block_start: &Hash256) -> Result<Hash256, ForkChoiceError> {
+        let start = self
+            .block_store
+            .get_deserialized(&justified_block_start)?
+            .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(*justified_block_start))?;
+
+        let start_state_root = start.state_root();
+
+        let latest_votes = self.get_latest_votes(&start_state_root, start.slot())?;
+
+        let mut head_hash = Hash256::zero();
+
+        loop {
+            let mut head_vote_count = 0;
+
+            let children = match self.children.get(&head_hash) {
+                Some(children) => children,
+                // we have found the head, exit
+                None => break,
+            };
+
+            for child_hash in children {
+                let vote_count = self.get_vote_count(&latest_votes, &child_hash)?;
+
+                if vote_count > head_vote_count {
+                    head_hash = *child_hash;
+                    head_vote_count = vote_count;
+                }
+            }
+        }
+        Ok(head_hash)
+    }
+}

eth2/types/src/lib.rs

@@ -24,7 +24,7 @@ pub mod readers;
 pub mod shard_reassignment_record;
 pub mod slashable_attestation;
 pub mod slashable_vote_data;
-pub mod slot_epoch;
+pub mod slot_epoch_height;
 pub mod spec;
 pub mod validator;
 pub mod validator_registry;
@@ -55,7 +55,7 @@ pub use crate::proposal_signed_data::ProposalSignedData;
 pub use crate::proposer_slashing::ProposerSlashing;
 pub use crate::slashable_attestation::SlashableAttestation;
 pub use crate::slashable_vote_data::SlashableVoteData;
-pub use crate::slot_epoch::{Epoch, Slot};
+pub use crate::slot_epoch_height::{Epoch, Slot};
 pub use crate::spec::ChainSpec;
 pub use crate::validator::{StatusFlags as ValidatorStatusFlags, Validator};
 pub use crate::validator_registry_delta_block::ValidatorRegistryDeltaBlock;

eth2/types/src/slot_epoch_height.rs

@@ -1,14 +1,14 @@
-/// The `Slot` and `Epoch` types are defined as newtypes over u64 to enforce type-safety between
+/// The `Slot` `Epoch`, `Height` types are defined as newtypes over u64 to enforce type-safety between
-/// the two types.
+/// the three types.
 ///
-/// `Slot` and `Epoch` have implementations which permit conversion, comparison and math operations
+/// `Slot`, `Epoch` and `Height` have implementations which permit conversion, comparison and math operations
 /// between each and `u64`, however specifically not between each other.
 ///
 /// All math operations on `Slot` and `Epoch` are saturating, they never wrap.
 ///
 /// It would be easy to define `PartialOrd` and other traits generically across all types which
-/// implement `Into<u64>`, however this would allow operations between `Slots` and `Epochs` which
+/// implement `Into<u64>`, however this would allow operations between `Slots`, `Epochs` and
-/// may lead to programming errors which are not detected by the compiler.
+/// `Heights` which may lead to programming errors which are not detected by the compiler.
 use crate::test_utils::TestRandom;
 use rand::RngCore;
 use serde_derive::Serialize;
@@ -42,6 +42,23 @@ macro_rules! impl_from_into_u64 {
     };
 }
 
+// need to truncate for some fork-choice algorithms
+macro_rules! impl_into_u32 {
+    ($main: ident) => {
+        impl Into<u32> for $main {
+            fn into(self) -> u32 {
+                self.0 as u32
+            }
+        }
+
+        impl $main {
+            pub fn as_u32(&self) -> u32 {
+                self.0 as u32
+            }
+        }
+    };
+}
+
 macro_rules! impl_from_into_usize {
     ($main: ident) => {
         impl From<usize> for $main {
@@ -269,13 +286,21 @@ macro_rules! impl_common {
     };
 }
 
+/// Beacon block slot.
 #[derive(Eq, Debug, Clone, Copy, Default, Serialize)]
 pub struct Slot(u64);
 
+/// Beacon block height, effectively `Slot/GENESIS_START_BLOCK`.
+#[derive(Eq, Debug, Clone, Copy, Default, Serialize)]
+pub struct Height(u64);
+
+/// Beacon Epoch, effectively `Slot / EPOCH_LENGTH`.
 #[derive(Eq, Debug, Clone, Copy, Default, Serialize)]
 pub struct Epoch(u64);
 
 impl_common!(Slot);
+impl_common!(Height);
+impl_into_u32!(Height); // height can be converted to u32
 impl_common!(Epoch);
 
 impl Slot {
@@ -287,11 +312,33 @@ impl Slot {
         Epoch::from(self.0 / epoch_length)
     }
 
+    pub fn height(self, genesis_slot: Slot) -> Height {
+        Height::from(self.0.saturating_sub(genesis_slot.as_u64()))
+    }
+
     pub fn max_value() -> Slot {
         Slot(u64::max_value())
     }
 }
 
+impl Height {
+    pub fn new(slot: u64) -> Height {
+        Height(slot)
+    }
+
+    pub fn slot(self, genesis_slot: Slot) -> Slot {
+        Slot::from(self.0.saturating_add(genesis_slot.as_u64()))
+    }
+
+    pub fn epoch(self, genesis_slot: u64, epoch_length: u64) -> Epoch {
+        Epoch::from(self.0.saturating_add(genesis_slot) / epoch_length)
+    }
+
+    pub fn max_value() -> Height {
+        Height(u64::max_value())
+    }
+}
+
 impl Epoch {
     pub fn new(slot: u64) -> Epoch {
         Epoch(slot)