Integrate fork choice into beacon_chain.

- Adds fork_choice to beacon_chain struct.
- Adds add_attestation inside process_free_attestation.
- Adds add_block inside process_block.
- Shifts core fork-choice logic into lib.rs.

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -2,6 +2,7 @@ use db::{
     stores::{BeaconBlockStore, BeaconStateStore},
     ClientDB, DBError,
 };
+use fork_choice::{ForkChoice, ForkChoiceError};
 use genesis::{genesis_beacon_block, genesis_beacon_state};
 use log::{debug, trace};
 use parking_lot::{RwLock, RwLockReadGuard};
@@ -14,11 +15,8 @@ use types::{
     AttestationData, BeaconBlock, BeaconBlockBody, BeaconState, ChainSpec, Eth1Data,
     FreeAttestation, Hash256, PublicKey, Signature,
 };
-use fork_choice::{longest_chain, basic_lmd_ghost};
-use fork_choice::{ForkChoice};
 
 use crate::attestation_aggregator::{AttestationAggregator, Outcome as AggregationOutcome};
-use crate::attestation_targets::AttestationTargets;
 use crate::block_graph::BlockGraph;
 use crate::checkpoint::CheckPoint;
 
@@ -29,6 +27,9 @@ pub enum Error {
     CommitteesError(CommitteesError),
     DBInconsistent(String),
     DBError(String),
+    ForkChoiceError(ForkChoiceError),
+    MissingBeaconBlock(Hash256),
+    MissingBeaconState(Hash256),
 }
 
 #[derive(Debug, PartialEq)]
@@ -60,7 +61,7 @@ pub enum BlockProcessingOutcome {
     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,
@@ -70,14 +71,15 @@ pub struct BeaconChain<T: ClientDB + Sized, U: SlotClock> {
     finalized_head: RwLock<CheckPoint>,
     justified_head: RwLock<CheckPoint>,
     pub state: RwLock<BeaconState>,
-    pub latest_attestation_targets: RwLock<AttestationTargets>,
     pub spec: ChainSpec,
+    pub fork_choice: 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(
@@ -85,6 +87,7 @@ where
         block_store: Arc<BeaconBlockStore<T>>,
         slot_clock: U,
         spec: ChainSpec,
+        fork_choice: F,
     ) -> Result<Self, Error> {
         if spec.initial_validators.is_empty() {
             return Err(Error::InsufficientValidators);
@@ -121,8 +124,6 @@ where
         ));
         let attestation_aggregator = RwLock::new(AttestationAggregator::new());
 
-        let latest_attestation_targets = RwLock::new(AttestationTargets::new());
-
         Ok(Self {
             block_store,
             state_store,
@@ -133,8 +134,8 @@ where
             justified_head,
             finalized_head,
             canonical_head,
-            latest_attestation_targets,
             spec: spec,
+            fork_choice,
         })
     }
 
@@ -209,7 +210,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> {
@@ -338,29 +339,27 @@ where
     /// - Create a new `Attestation`.
     /// - Aggregate it to an existing `Attestation`.
     pub fn process_free_attestation(
-        &self,
+        &mut 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.clone()),
-            None => None,
         }
+
+        // valid attestation, proceed with fork-choice logic
+        self.fork_choice.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.
@@ -417,7 +416,7 @@ where
     /// Accept some block and attempt to add it to block DAG.
     ///
     /// Will accept blocks from prior slots, however it will reject any block from a future slot.
-    pub fn process_block(&self, block: BeaconBlock) -> Result<BlockProcessingOutcome, Error> {
+    pub fn process_block(&mut self, block: BeaconBlock) -> Result<BlockProcessingOutcome, Error> {
         debug!("Processing block with slot {}...", block.slot());
 
         let block_root = block.canonical_root();
@@ -471,7 +470,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(
@@ -495,6 +494,9 @@ where
         self.block_graph
             .add_leaf(&parent_block_root, block_root.clone());
 
+        // run the fork_choice add_block logic
+        self.fork_choice.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
@@ -574,15 +576,11 @@ where
         Some((block, state))
     }
 
-    // For now, we give it the option of choosing which fork choice to use
+    // TODO: Left this as is, modify later
-    pub fn fork_choice(&self, fork_choice: ForkChoice) -> Result<(), Error> {
+    pub fn fork_choice(&mut self) -> Result<(), Error> {
-        let present_head = &self.finalized_head().beacon_block_root;
+        let present_head = &self.finalized_head().beacon_block_root.clone();
 
-        let new_head = match fork_choice {
+        let new_head = self.fork_choice.find_head(present_head)?;
-           ForkChoice::BasicLMDGhost => basic_lmd_ghost(&self.finalized_head().beacon_block_root)?,
-           // TODO: Implement others
-           _ => present_head
-        }
 
         if new_head != *present_head {
             let block = self
@@ -601,18 +599,21 @@ where
         }
 
         Ok(())
-
+    }
-
-
 }
 
-
 impl From<DBError> for Error {
     fn from(e: DBError) -> Error {
         Error::DBError(e.message)
     }
 }
 
+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/lib.rs

@@ -1,5 +1,4 @@
 mod attestation_aggregator;
-mod attestation_targets;
 mod beacon_chain;
 mod block_graph;
 mod checkpoint;

eth2/fork_choice/src/lib.rs

@@ -23,8 +23,51 @@ pub mod longest_chain;
 pub mod optimised_lmd_ghost;
 pub mod protolambda_lmd_ghost;
 
+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 {
+    /// 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)
+    }
+}
+
 /// Fork choice options that are currently implemented.
-pub enum ForkChoice {
+pub enum ForkChoiceAlgorithms {
     /// Chooses the longest chain becomes the head. Not for production.
     LongestChain,
     /// A simple and highly inefficient implementation of LMD ghost.

eth2/fork_choice/src/optimised_lmd_ghost.rs

@@ -1,9 +1,10 @@
 extern crate byteorder;
 extern crate fast_math;
+use crate::{ForkChoice, ForkChoiceError};
 use byteorder::{BigEndian, ByteOrder};
 use db::{
     stores::{BeaconBlockStore, BeaconStateStore},
-    ClientDB, DBError,
+    ClientDB,
 };
 use fast_math::log2_raw;
 use std::collections::HashMap;
@@ -11,7 +12,7 @@ use std::sync::Arc;
 use types::{
     readers::{BeaconBlockReader, BeaconStateReader},
     validator_registry::get_active_validator_indices,
-    Attestation, BeaconBlock, Hash256,
+    BeaconBlock, Hash256,
 };
 
 //TODO: Sort out global constants
@@ -46,7 +47,7 @@ pub struct OptimisedLMDGhost<T: ClientDB + Sized> {
     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<usize, Hash256>,
+    latest_attestation_targets: HashMap<u64, Hash256>,
     /// Block storage access.
     block_store: Arc<BeaconBlockStore<T>>,
     /// State storage access.
@@ -82,7 +83,7 @@ where
                 .into_beacon_block()?
                 .slot;
 
-            (block_slot - self.GENESIS_SLOT) as u32
+            (block_slot - GENESIS_SLOT) as u32
         };
 
         // verify we haven't exceeded the block height
@@ -202,7 +203,7 @@ impl<T: ClientDB + Sized> ForkChoice for OptimisedLMDGhost<T> {
             .get_reader(&block.parent_root)?
             .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(block.parent_root))?
             .slot()
-            - self.GENESIS_SLOT;
+            - GENESIS_SLOT;
 
         let parent_hash = &block.parent_root;
 
@@ -228,8 +229,14 @@ impl<T: ClientDB + Sized> ForkChoice for OptimisedLMDGhost<T> {
         Ok(())
     }
 
-    fn add_attestation(&mut self, attestation: &Attestation) -> Result<(), ForkChoiceError> {
+    fn add_attestation(
-        // simply add the attestation to the latest_message_target mapping
+        &mut self,
+        validator_index: u64,
+        target_block_root: &Hash256,
+    ) -> Result<(), ForkChoiceError> {
+        // simply add the attestation to the latest_attestation_target
+        self.latest_attestation_targets
+            .insert(validator_index, target_block_root.clone());
         Ok(())
     }
 
@@ -242,7 +249,7 @@ impl<T: ClientDB + Sized> ForkChoice for OptimisedLMDGhost<T> {
         //.into_beacon_block()?;
 
         let block_slot = block.slot();
-        let block_height = block_slot - self.GENESIS_SLOT;
+        let block_height = block_slot - GENESIS_SLOT;
         let state_root = block.state_root();
 
         // get latest votes
@@ -262,12 +269,12 @@ impl<T: ClientDB + Sized> ForkChoice for OptimisedLMDGhost<T> {
             let active_validator_indices =
                 get_active_validator_indices(&current_state.validator_registry, block_slot);
 
-            for i in active_validator_indices {
+            for index in active_validator_indices {
                 let balance =
-                    std::cmp::min(current_state.validator_balances[i], self.MAX_DEPOSIT_AMOUNT)
+                    std::cmp::min(current_state.validator_balances[index], MAX_DEPOSIT_AMOUNT)
-                        / self.FORK_CHOICE_BALANCE_INCREMENT;
+                        / FORK_CHOICE_BALANCE_INCREMENT;
                 if balance > 0 {
-                    if let Some(target) = self.latest_attestation_targets.get(&(i as usize)) {
+                    if let Some(target) = self.latest_attestation_targets.get(&(index as u64)) {
                         *latest_votes.entry(*target).or_insert_with(|| 0) += balance;
                     }
                 }
@@ -331,7 +338,7 @@ impl<T: ClientDB + Sized> ForkChoice for OptimisedLMDGhost<T> {
                 .get_reader(&current_head)?
                 .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(*justified_block_start))?
                 .slot()
-                - self.GENESIS_SLOT;
+                - 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
@@ -342,40 +349,6 @@ impl<T: ClientDB + Sized> ForkChoice for OptimisedLMDGhost<T> {
     }
 }
 
-/// 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 {
-    /// 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.
-    fn add_attestation(&mut self, attestation: &Attestation) -> 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.
-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)
-    }
-}
-
 /// 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> {