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,52 +1,73 @@
-use db::stores::BeaconBlockStore;
+use crate::{ForkChoice, ForkChoiceError};
-use db::{ClientDB, DBError};
+use db::{stores::BeaconBlockStore, ClientDB};
-use ssz::{Decodable, DecodeError};
 use std::sync::Arc;
 use types::{BeaconBlock, Hash256, Slot};
 
-pub enum ForkChoiceError {
+pub struct LongestChain<T>
-    BadSszInDatabase,
-    MissingBlock,
-    DBError(String),
-}
-
-pub fn longest_chain<T>(
-    head_block_hashes: &[Hash256],
-    block_store: &Arc<BeaconBlockStore<T>>,
-) -> Result<Option<usize>, ForkChoiceError>
 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>>,
+}
 
+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<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(())
+    }
+
+    fn add_attestation(&mut self, _: u64, _: &Hash256) -> Result<(), ForkChoiceError> {
+        // do nothing
+        Ok(())
+    }
+
+    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 head_block_hashes.iter().enumerate() {
+        for (index, block_hash) in self.head_block_hashes.iter().enumerate() {
-        let ssz = block_store
+            let block = self
-            .get(&block_hash)?
+                .block_store
-            .ok_or(ForkChoiceError::MissingBlock)?;
+                .get_deserialized(&block_hash)?
-        let (block, _) = BeaconBlock::ssz_decode(&ssz, 0)?;
+                .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: Option<Slot> = None;
+        let highest_slot = head_blocks
-    for (_, block) in &head_blocks {
+            .iter()
-        let slot = block.slot;
+            .fold(Slot::from(0u64), |highest, (_, block)| {
+                std::cmp::max(block.slot, highest)
+            });
 
-        match highest_slot {
+        // if we find no blocks, return Error
-            None => Some(slot),
+        if highest_slot == 0 {
-            Some(winning_slot) => {
+            return Err(ForkChoiceError::HeadNotFound);
-                if slot > winning_slot {
-                    Some(slot)
-                } else {
-                    Some(winning_slot)
-                }
-            }
-        };
         }
 
         /*
@@ -55,39 +76,27 @@ where
          * 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),
+        let head_index: Option<usize> =
-        Some(highest_slot) => {
+            head_blocks
-            let mut highest_blocks = vec![];
+                .iter()
-            for (index, block) in head_blocks {
+                .fold(None, |smallest_index, (index, block)| {
                     if block.slot == highest_slot {
-                    highest_blocks.push((index, block))
+                        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);
         }
 
-            highest_blocks.sort_by(|a, b| head_block_hashes[a.0].cmp(&head_block_hashes[b.0]));
+        Ok(self.head_block_hashes[head_index.unwrap()])
-            let (index, _) = highest_blocks[0];
-            Ok(Some(index))
-        }
-    }
-}
-
-impl From<DecodeError> for ForkChoiceError {
-    fn from(_: DecodeError) -> Self {
-        ForkChoiceError::BadSszInDatabase
-    }
-}
-
-impl From<DBError> for ForkChoiceError {
-    fn from(e: DBError) -> Self {
-        ForkChoiceError::DBError(e.message)
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    #[test]
-    fn test_naive_fork_choice() {
-        assert_eq!(2 + 2, 4);
     }
 }