Update fork choice find head fn

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -74,7 +74,7 @@ pub struct BeaconChain<T: BeaconChainTypes> {
     genesis_block_root: Hash256,
     /// A state-machine that is updated with information from the network and chooses a canonical
     /// head block.
-    pub fork_choice: ForkChoice<T::LmdGhost, T::Store, T::EthSpec>,
+    pub fork_choice: ForkChoice<T>,
     /// Stores metrics about this `BeaconChain`.
     pub metrics: Metrics,
 }
@@ -87,7 +87,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         mut genesis_state: BeaconState<T::EthSpec>,
         genesis_block: BeaconBlock,
         spec: ChainSpec,
-        fork_choice: ForkChoice<T::LmdGhost, T::Store, T::EthSpec>,
+        fork_choice: ForkChoice<T>,
     ) -> Result<Self, Error> {
         let state_root = genesis_state.canonical_root();
         store.put(&state_root, &genesis_state)?;
@@ -714,7 +714,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         let timer = self.metrics.fork_choice_times.start_timer();
 
         // Determine the root of the block that is the head of the chain.
-        let beacon_block_root = self.fork_choice.find_head()?;
+        let beacon_block_root = self.fork_choice.find_head(&self)?;
 
         // End fork choice metrics timer.
         timer.observe_duration();

beacon_node/beacon_chain/src/fork_choice.rs

@@ -1,46 +1,75 @@
-use crate::BeaconChain;
+use crate::{BeaconChain, BeaconChainTypes};
 use lmd_ghost::LmdGhost;
 use state_processing::common::get_attesting_indices_unsorted;
-use std::marker::PhantomData;
 use std::sync::Arc;
-use store::Store;
+use store::{Error as StoreError, Store};
 use types::{Attestation, BeaconBlock, BeaconState, BeaconStateError, EthSpec, Hash256};
 
 type Result<T> = std::result::Result<T, Error>;
 
 #[derive(Debug, PartialEq)]
 pub enum Error {
+    MissingBlock(Hash256),
+    MissingState(Hash256),
     BackendError(String),
     BeaconStateError(BeaconStateError),
+    StoreError(StoreError),
 }
 
-pub struct ForkChoice<L, S, E> {
-    backend: L,
-    _phantom_a: PhantomData<S>,
-    _phantom_b: PhantomData<E>,
+pub struct ForkChoice<T: BeaconChainTypes> {
+    backend: T::LmdGhost,
 }
 
-impl<L, S, E> ForkChoice<L, S, E>
-where
-    L: LmdGhost<S, E>,
-    S: Store,
-    E: EthSpec,
-{
-    pub fn new(store: Arc<S>) -> Self {
+impl<T: BeaconChainTypes> ForkChoice<T> {
+    pub fn new(store: Arc<T::Store>) -> Self {
         Self {
-            backend: L::new(store),
-            _phantom_a: PhantomData,
-            _phantom_b: PhantomData,
+            backend: T::LmdGhost::new(store),
         }
     }
 
-    pub fn find_head(&self) -> Result<Hash256> {
-        self.backend.find_head().map_err(Into::into)
+    pub fn find_head(&self, chain: &BeaconChain<T>) -> Result<Hash256> {
+        // From the specification:
+        //
+        // Let justified_head be the descendant of finalized_head with the highest epoch that has
+        // been justified for at least 1 epoch ... If no such descendant exists,
+        // set justified_head to finalized_head.
+        let (start_state, start_block_root) = {
+            let state = chain.current_state();
+
+            let block_root = if state.current_epoch() + 1 > state.current_justified_epoch {
+                state.current_justified_root
+            } else {
+                state.finalized_root
+            };
+            let block = chain
+                .store
+                .get::<BeaconBlock>(&block_root)?
+                .ok_or_else(|| Error::MissingBlock(block_root))?;
+
+            let state = chain
+                .store
+                .get::<BeaconState<T::EthSpec>>(&block.state_root)?
+                .ok_or_else(|| Error::MissingState(block.state_root))?;
+
+            (state, block_root)
+        };
+
+        // A function that returns the weight for some validator index.
+        let weight = |validator_index: usize| -> Option<u64> {
+            start_state
+                .validator_registry
+                .get(validator_index)
+                .and_then(|v| Some(v.effective_balance))
+        };
+
+        self.backend
+            .find_head(start_block_root, weight)
+            .map_err(Into::into)
     }
 
     pub fn process_attestation(
         &self,
-        state: &BeaconState<E>,
+        state: &BeaconState<T::EthSpec>,
         attestation: &Attestation,
     ) -> Result<()> {
         // Note: `get_attesting_indices_unsorted` requires that the beacon state caches be built.
@@ -56,7 +85,7 @@ where
         let block_slot = attestation
             .data
             .target_epoch
-            .start_slot(E::slots_per_epoch());
+            .start_slot(T::EthSpec::slots_per_epoch());
 
         for validator_index in validator_indices {
             self.backend
@@ -70,7 +99,11 @@ where
     ///
     /// Assumes the block (and therefore it's attestations) are valid. It is a logic error to
     /// provide an invalid block.
-    pub fn process_block(&self, state: &BeaconState<E>, block: &BeaconBlock) -> Result<()> {
+    pub fn process_block(
+        &self,
+        state: &BeaconState<T::EthSpec>,
+        block: &BeaconBlock,
+    ) -> Result<()> {
         // Note: we never count the block as a latest message, only attestations.
         //
         // I (Paul H) do not have an explicit reference to this, however I derive it from this
@@ -91,6 +124,12 @@ impl From<BeaconStateError> for Error {
     }
 }
 
+impl From<StoreError> for Error {
+    fn from(e: StoreError) -> Error {
+        Error::StoreError(e)
+    }
+}
+
 impl From<String> for Error {
     fn from(e: String) -> Error {
         Error::BackendError(e)

eth2/lmd_ghost/src/lib.rs

@@ -18,5 +18,7 @@ pub trait LmdGhost<S: Store, E: EthSpec>: Send + Sync {
         block_slot: Slot,
     ) -> Result<()>;
 
-    fn find_head(&self) -> Result<Hash256>;
+    fn find_head<F>(&self, start_block_root: Hash256, weight: F) -> Result<Hash256>
+    where
+        F: Fn(usize) -> Option<u64>;
 }

eth2/lmd_ghost/src/reduced_tree.rs

@@ -86,7 +86,10 @@ where
             .map_err(Into::into)
     }
 
-    fn find_head(&self) -> SuperResult<Hash256> {
+    fn find_head<F>(&self, _start_block_root: Hash256, _weight: F) -> SuperResult<Hash256>
+    where
+        F: Fn(usize) -> Option<u64>,
+    {
         unimplemented!();
     }
 }