diff --git a/eth2/fork_choice/src/lib.rs b/eth2/fork_choice/src/lib.rs index 820a33ec1d..a6e4ea422e 100644 --- a/eth2/fork_choice/src/lib.rs +++ b/eth2/fork_choice/src/lib.rs @@ -83,6 +83,7 @@ pub enum ForkChoiceError { CannotFindBestChild, ChildrenNotFound, StorageError(String), + HeadNotFound, } impl From for ForkChoiceError { diff --git a/eth2/fork_choice/src/longest_chain.rs b/eth2/fork_choice/src/longest_chain.rs index 277d6b9509..c25cd43a17 100644 --- a/eth2/fork_choice/src/longest_chain.rs +++ b/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( - head_block_hashes: &[Hash256], - block_store: &Arc>, -) -> Result, ForkChoiceError> +pub struct LongestChain where T: ClientDB + Sized, { - let mut head_blocks: Vec<(usize, BeaconBlock)> = vec![]; + /// List of head block hashes + head_block_hashes: Vec, + /// Block storage access. + block_store: Arc>, +} - /* - * 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 = 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 LongestChain +where + T: ClientDB + Sized, +{ + pub fn new(block_store: Arc>) -> Self { + LongestChain { + head_block_hashes: Vec::new(), + block_store, } } } -impl From for ForkChoiceError { - fn from(_: DecodeError) -> Self { - ForkChoiceError::BadSszInDatabase +impl ForkChoice for LongestChain { + 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 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 { + 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 = + 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()]) } }