Updates longest chain to match new fork-choice structure.

eth2/fork_choice/src/lib.rs

@@ -83,6 +83,7 @@ pub enum ForkChoiceError {
     CannotFindBestChild,
     ChildrenNotFound,
     StorageError(String),
+    HeadNotFound,
 }
 
 impl From<DBError> for ForkChoiceError {

eth2/fork_choice/src/longest_chain.rs

@@ -1,93 +1,102 @@
-use db::stores::BeaconBlockStore;
-use db::{ClientDB, DBError};
-use ssz::{Decodable, DecodeError};
+use crate::{ForkChoice, ForkChoiceError};
+use db::{stores::BeaconBlockStore, ClientDB};
 use std::sync::Arc;
 use types::{BeaconBlock, Hash256, Slot};
 
-pub enum ForkChoiceError {
-    BadSszInDatabase,
-    MissingBlock,
-    DBError(String),
-}
-
-pub fn longest_chain<T>(
-    head_block_hashes: &[Hash256],
-    block_store: &Arc<BeaconBlockStore<T>>,
-) -> Result<Option<usize>, ForkChoiceError>
+pub struct LongestChain<T>
 where
     T: ClientDB + Sized,
 {
-    let mut head_blocks: Vec<(usize, BeaconBlock)> = vec![];
+    /// List of head block hashes
+    head_block_hashes: Vec<Hash256>,
+    /// Block storage access.
+    block_store: Arc<BeaconBlockStore<T>>,
+}
 
-    /*
-     * Load all the head_block hashes from the DB as SszBeaconBlocks.
-     */
-    for (index, block_hash) in head_block_hashes.iter().enumerate() {
-        let ssz = block_store
-            .get(&block_hash)?
-            .ok_or(ForkChoiceError::MissingBlock)?;
-        let (block, _) = BeaconBlock::ssz_decode(&ssz, 0)?;
-        head_blocks.push((index, block));
-    }
-
-    /*
-     * Loop through all the head blocks and find the highest slot.
-     */
-    let highest_slot: Option<Slot> = None;
-    for (_, block) in &head_blocks {
-        let slot = block.slot;
-
-        match highest_slot {
-            None => Some(slot),
-            Some(winning_slot) => {
-                if slot > winning_slot {
-                    Some(slot)
-                } else {
-                    Some(winning_slot)
-                }
-            }
-        };
-    }
-
-    /*
-     * Loop through all the highest blocks and sort them by highest hash.
-     *
-     * Ultimately, the index of the head_block hash with the highest slot and highest block
-     * hash will be the winner.
-     */
-    match highest_slot {
-        None => Ok(None),
-        Some(highest_slot) => {
-            let mut highest_blocks = vec![];
-            for (index, block) in head_blocks {
-                if block.slot == highest_slot {
-                    highest_blocks.push((index, block))
-                }
-            }
-
-            highest_blocks.sort_by(|a, b| head_block_hashes[a.0].cmp(&head_block_hashes[b.0]));
-            let (index, _) = highest_blocks[0];
-            Ok(Some(index))
+impl<T> LongestChain<T>
+where
+    T: ClientDB + Sized,
+{
+    pub fn new(block_store: Arc<BeaconBlockStore<T>>) -> Self {
+        LongestChain {
+            head_block_hashes: Vec::new(),
+            block_store,
         }
     }
 }
 
-impl From<DecodeError> for ForkChoiceError {
-    fn from(_: DecodeError) -> Self {
-        ForkChoiceError::BadSszInDatabase
+impl<T: ClientDB + Sized> ForkChoice for LongestChain<T> {
+    fn add_block(
+        &mut self,
+        block: &BeaconBlock,
+        block_hash: &Hash256,
+    ) -> Result<(), ForkChoiceError> {
+        // add the block hash to head_block_hashes removing the parent if it exists
+        self.head_block_hashes
+            .retain(|hash| *hash != block.parent_root);
+        self.head_block_hashes.push(*block_hash);
+        Ok(())
     }
-}
 
-impl From<DBError> for ForkChoiceError {
-    fn from(e: DBError) -> Self {
-        ForkChoiceError::DBError(e.message)
+    fn add_attestation(&mut self, _: u64, _: &Hash256) -> Result<(), ForkChoiceError> {
+        // do nothing
+        Ok(())
     }
-}
 
-#[cfg(test)]
-mod tests {
-    #[test]
-    fn test_naive_fork_choice() {
-        assert_eq!(2 + 2, 4);
+    fn find_head(&mut self, _: &Hash256) -> Result<Hash256, ForkChoiceError> {
+        let mut head_blocks: Vec<(usize, BeaconBlock)> = vec![];
+        /*
+         * Load all the head_block hashes from the DB as SszBeaconBlocks.
+         */
+        for (index, block_hash) in self.head_block_hashes.iter().enumerate() {
+            let block = self
+                .block_store
+                .get_deserialized(&block_hash)?
+                .ok_or(ForkChoiceError::MissingBeaconBlock(*block_hash))?;
+            head_blocks.push((index, block));
+        }
+
+        /*
+         * Loop through all the head blocks and find the highest slot.
+         */
+        let highest_slot = head_blocks
+            .iter()
+            .fold(Slot::from(0u64), |highest, (_, block)| {
+                std::cmp::max(block.slot, highest)
+            });
+
+        // if we find no blocks, return Error
+        if highest_slot == 0 {
+            return Err(ForkChoiceError::HeadNotFound);
+        }
+
+        /*
+         * Loop through all the highest blocks and sort them by highest hash.
+         *
+         * Ultimately, the index of the head_block hash with the highest slot and highest block
+         * hash will be the winner.
+         */
+
+        let head_index: Option<usize> =
+            head_blocks
+                .iter()
+                .fold(None, |smallest_index, (index, block)| {
+                    if block.slot == highest_slot {
+                        if smallest_index.is_none() {
+                            return Some(*index);
+                        }
+                        return Some(std::cmp::min(
+                            *index,
+                            smallest_index.expect("Cannot be None"),
+                        ));
+                    }
+                    smallest_index
+                });
+
+        if head_index.is_none() {
+            return Err(ForkChoiceError::HeadNotFound);
+        }
+
+        Ok(self.head_block_hashes[head_index.unwrap()])
     }
 }