diff --git a/eth2/fork_choice/src/lib.rs b/eth2/fork_choice/src/lib.rs index 6062c19b1..0d6969e89 100644 --- a/eth2/fork_choice/src/lib.rs +++ b/eth2/fork_choice/src/lib.rs @@ -22,6 +22,7 @@ extern crate types; pub mod bitwise_lmd_ghost; pub mod longest_chain; +pub mod optimized_lmd_ghost; pub mod slow_lmd_ghost; use db::stores::BeaconBlockAtSlotError; @@ -30,6 +31,7 @@ use types::{BeaconBlock, ChainSpec, Hash256}; pub use bitwise_lmd_ghost::BitwiseLMDGhost; pub use longest_chain::LongestChain; +pub use optimized_lmd_ghost::OptimizedLMDGhost; pub use slow_lmd_ghost::SlowLMDGhost; /// Defines the interface for Fork Choices. Each Fork choice will define their own data structures @@ -101,4 +103,6 @@ pub enum ForkChoiceAlgorithm { SlowLMDGhost, /// An optimised version of bitwise LMD-GHOST by Vitalik. BitwiseLMDGhost, + /// An optimised implementation of LMD ghost. + OptimizedLMDGhost, } diff --git a/eth2/fork_choice/src/optimized_lmd_ghost.rs b/eth2/fork_choice/src/optimized_lmd_ghost.rs new file mode 100644 index 000000000..636ccdabc --- /dev/null +++ b/eth2/fork_choice/src/optimized_lmd_ghost.rs @@ -0,0 +1,453 @@ +//! The optimised bitwise LMD-GHOST fork choice rule. +extern crate bit_vec; + +use crate::{ForkChoice, ForkChoiceError}; +use db::{ + stores::{BeaconBlockStore, BeaconStateStore}, + ClientDB, +}; +use log::{debug, trace}; +use std::cmp::Ordering; +use std::collections::HashMap; +use std::sync::Arc; +use types::{ + readers::BeaconBlockReader, validator_registry::get_active_validator_indices, BeaconBlock, + ChainSpec, Hash256, Slot, SlotHeight, +}; + +//TODO: Pruning - Children +//TODO: Handle Syncing + +// NOTE: This uses u32 to represent difference between block heights. Thus this is only +// applicable for block height differences in the range of a u32. +// This can potentially be parallelized in some parts. + +/// Compute the base-2 logarithm of an integer, floored (rounded down) +#[inline] +fn log2_int(x: u64) -> u32 { + if x == 0 { + return 0; + } + 63 - x.leading_zeros() +} + +fn power_of_2_below(x: u64) -> u64 { + 2u64.pow(log2_int(x)) +} + +/// Stores the necessary data structures to run the optimised lmd ghost algorithm. +pub struct OptimizedLMDGhost { + /// A cache of known ancestors at given heights for a specific block. + //TODO: Consider FnvHashMap + cache: HashMap, Hash256>, + /// Log lookup table for blocks to their ancestors. + //TODO: Verify we only want/need a size 16 log lookup + ancestors: Vec>, + /// Stores the children for any given parent. + children: HashMap>, + /// 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, + /// Block storage access. + block_store: Arc>, + /// State storage access. + state_store: Arc>, + max_known_height: SlotHeight, +} + +impl OptimizedLMDGhost +where + T: ClientDB + Sized, +{ + pub fn new( + block_store: Arc>, + state_store: Arc>, + ) -> Self { + OptimizedLMDGhost { + cache: HashMap::new(), + ancestors: vec![HashMap::new(); 16], + latest_attestation_targets: HashMap::new(), + children: HashMap::new(), + max_known_height: SlotHeight::new(0), + block_store, + state_store, + } + } + + /// Finds the latest votes weighted by validator balance. Returns a hashmap of block_hash to + /// weighted votes. + pub fn get_latest_votes( + &self, + state_root: &Hash256, + block_slot: Slot, + spec: &ChainSpec, + ) -> Result, ForkChoiceError> { + // get latest votes + // Note: Votes are weighted by min(balance, MAX_DEPOSIT_AMOUNT) // + // FORK_CHOICE_BALANCE_INCREMENT + // build a hashmap of block_hash to weighted votes + let mut latest_votes: HashMap = HashMap::new(); + // gets the current weighted votes + let current_state = self + .state_store + .get_deserialized(&state_root)? + .ok_or_else(|| ForkChoiceError::MissingBeaconState(*state_root))?; + + let active_validator_indices = get_active_validator_indices( + ¤t_state.validator_registry[..], + block_slot.epoch(spec.slots_per_epoch), + ); + + for index in active_validator_indices { + let balance = std::cmp::min( + current_state.validator_balances[index], + spec.max_deposit_amount, + ) / spec.fork_choice_balance_increment; + if balance > 0 { + if let Some(target) = self.latest_attestation_targets.get(&(index as u64)) { + *latest_votes.entry(*target).or_insert_with(|| 0) += balance; + } + } + } + trace!("Latest votes: {:?}", latest_votes); + Ok(latest_votes) + } + + /// Gets the ancestor at a given height `at_height` of a block specified by `block_hash`. + fn get_ancestor( + &mut self, + block_hash: Hash256, + target_height: SlotHeight, + spec: &ChainSpec, + ) -> Option { + // return None if we can't get the block from the db. + let block_height = { + let block_slot = self + .block_store + .get_deserialized(&block_hash) + .ok()? + .expect("Should have returned already if None") + .slot; + + block_slot.height(spec.genesis_slot) + }; + + // verify we haven't exceeded the block height + if target_height >= block_height { + if target_height > block_height { + return None; + } else { + return Some(block_hash); + } + } + // check if the result is stored in our cache + let cache_key = CacheKey::new(&block_hash, target_height.as_u64()); + if let Some(ancestor) = self.cache.get(&cache_key) { + return Some(*ancestor); + } + + // not in the cache recursively search for ancestors using a log-lookup + if let Some(ancestor) = { + let ancestor_lookup = self.ancestors + [log2_int((block_height - target_height - 1u64).as_u64()) as usize] + .get(&block_hash) + //TODO: Panic if we can't lookup and fork choice fails + .expect("All blocks should be added to the ancestor log lookup table"); + self.get_ancestor(*ancestor_lookup, target_height, &spec) + } { + // add the result to the cache + self.cache.insert(cache_key, ancestor); + return Some(ancestor); + } + + None + } + + // looks for an obvious block winner given the latest votes for a specific height + fn get_clear_winner( + &mut self, + latest_votes: &HashMap, + block_height: SlotHeight, + spec: &ChainSpec, + ) -> Option { + // map of vote counts for every hash at this height + let mut current_votes: HashMap = HashMap::new(); + let mut total_vote_count = 0; + + trace!("Clear winner at block height: {}", block_height); + // loop through the latest votes and count all votes + // these have already been weighted by balance + for (hash, votes) in latest_votes.iter() { + if let Some(ancestor) = self.get_ancestor(*hash, block_height, spec) { + let current_vote_value = current_votes.get(&ancestor).unwrap_or_else(|| &0); + current_votes.insert(ancestor, current_vote_value + *votes); + total_vote_count += votes; + } + } + // Check if there is a clear block winner at this height. If so return it. + for (hash, votes) in current_votes.iter() { + if *votes > total_vote_count / 2 { + // we have a clear winner, return it + return Some(*hash); + } + } + // didn't find a clear winner + None + } + + // Finds the best child (one with highest votes) + fn choose_best_child(&self, votes: &HashMap) -> Option { + if votes.is_empty() { + return None; + } + + // Iterate through hashmap to get child with maximum votes + let best_child = votes.iter().max_by(|(child1, v1), (child2, v2)| { + let mut result = v1.cmp(v2); + // If votes are equal, choose smaller hash to break ties deterministically + if result == Ordering::Equal { + // Reverse so that max_by chooses smaller hash + result = child1.cmp(child2).reverse(); + } + result + }); + + Some(*best_child.unwrap().0) + } +} + +impl ForkChoice for OptimizedLMDGhost { + fn add_block( + &mut self, + block: &BeaconBlock, + block_hash: &Hash256, + spec: &ChainSpec, + ) -> Result<(), ForkChoiceError> { + // get the height of the parent + let parent_height = self + .block_store + .get_deserialized(&block.parent_root)? + .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(block.parent_root))? + .slot() + .height(spec.genesis_slot); + + let parent_hash = &block.parent_root; + + // add the new block to the children of parent + (*self + .children + .entry(block.parent_root) + .or_insert_with(|| vec![])) + .push(block_hash.clone()); + + // build the ancestor data structure + for index in 0..16 { + if parent_height % (1 << index) == 0 { + self.ancestors[index].insert(*block_hash, *parent_hash); + } else { + // TODO: This is unsafe. Will panic if parent_hash doesn't exist. Using it for debugging + let parent_ancestor = self.ancestors[index][parent_hash]; + self.ancestors[index].insert(*block_hash, parent_ancestor); + } + } + // update the max height + self.max_known_height = std::cmp::max(self.max_known_height, parent_height + 1); + Ok(()) + } + + fn add_attestation( + &mut self, + validator_index: u64, + target_block_root: &Hash256, + spec: &ChainSpec, + ) -> Result<(), ForkChoiceError> { + // simply add the attestation to the latest_attestation_target if the block_height is + // larger + trace!( + "Adding attestation of validator: {:?} for block: {}", + validator_index, + target_block_root + ); + let attestation_target = self + .latest_attestation_targets + .entry(validator_index) + .or_insert_with(|| *target_block_root); + // if we already have a value + if attestation_target != target_block_root { + trace!("Old attestation found: {:?}", attestation_target); + // get the height of the target block + let block_height = self + .block_store + .get_deserialized(&target_block_root)? + .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(*target_block_root))? + .slot() + .height(spec.genesis_slot); + + // get the height of the past target block + let past_block_height = self + .block_store + .get_deserialized(&attestation_target)? + .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(*attestation_target))? + .slot() + .height(spec.genesis_slot); + // update the attestation only if the new target is higher + if past_block_height < block_height { + trace!("Updating old attestation"); + *attestation_target = *target_block_root; + } + } + Ok(()) + } + + /// Perform lmd_ghost on the current chain to find the head. + fn find_head( + &mut self, + justified_block_start: &Hash256, + spec: &ChainSpec, + ) -> Result { + debug!( + "Starting optimised fork choice at block: {}", + justified_block_start + ); + let block = self + .block_store + .get_deserialized(&justified_block_start)? + .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(*justified_block_start))?; + + let block_slot = block.slot(); + let state_root = block.state_root(); + let mut block_height = block_slot.height(spec.genesis_slot); + + let mut current_head = *justified_block_start; + + let mut latest_votes = self.get_latest_votes(&state_root, block_slot, spec)?; + + // remove any votes that don't relate to our current head. + latest_votes + .retain(|hash, _| self.get_ancestor(*hash, block_height, spec) == Some(current_head)); + + // begin searching for the head + loop { + debug!( + "Iteration for block: {} with vote length: {}", + current_head, + latest_votes.len() + ); + // if there are no children, we are done, return the current_head + let children = match self.children.get(¤t_head) { + Some(children) => children.clone(), + None => { + debug!("Head found: {}", current_head); + return Ok(current_head); + } + }; + + // logarithmic lookup blocks to see if there are obvious winners, if so, + // progress to the next iteration. + let mut step = + power_of_2_below(self.max_known_height.saturating_sub(block_height).as_u64()) / 2; + while step > 0 { + trace!("Current Step: {}", step); + if let Some(clear_winner) = self.get_clear_winner( + &latest_votes, + block_height - (block_height % step) + step, + spec, + ) { + current_head = clear_winner; + break; + } + step /= 2; + } + if step > 0 { + trace!("Found clear winner: {}", current_head); + } + // if our skip lookup failed and we only have one child, progress to that child + else if children.len() == 1 { + current_head = children[0]; + trace!( + "Lookup failed, only one child, proceeding to child: {}", + current_head + ); + } + // we need to find the best child path to progress down. + else { + trace!("Searching for best child"); + let mut child_votes = HashMap::new(); + for (voted_hash, vote) in latest_votes.iter() { + // if the latest votes correspond to a child + if let Some(child) = self.get_ancestor(*voted_hash, block_height + 1, spec) { + // add up the votes for each child + *child_votes.entry(child).or_insert_with(|| 0) += vote; + } + } + // given the votes on the children, find the best child + current_head = self + .choose_best_child(&child_votes) + .ok_or(ForkChoiceError::CannotFindBestChild)?; + trace!("Best child found: {}", current_head); + } + + // didn't find head yet, proceed to next iteration + // update block height + block_height = self + .block_store + .get_deserialized(¤t_head)? + .ok_or_else(|| ForkChoiceError::MissingBeaconBlock(current_head))? + .slot() + .height(spec.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 + for hash in latest_votes.keys() { + trace!( + "Ancestor for vote: {} at height: {} is: {:?}", + hash, + block_height, + self.get_ancestor(*hash, block_height, spec) + ); + } + latest_votes.retain(|hash, _| { + self.get_ancestor(*hash, block_height, spec) == Some(current_head) + }); + } + } +} + +/// Type for storing blocks in a memory cache. Key is comprised of block-hash plus the height. +#[derive(PartialEq, Eq, Hash)] +pub struct CacheKey { + block_hash: Hash256, + block_height: T, +} + +impl CacheKey { + pub fn new(block_hash: &Hash256, block_height: T) -> Self { + CacheKey { + block_hash: *block_hash, + block_height, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + pub fn test_power_of_2_below() { + assert_eq!(power_of_2_below(4), 4); + assert_eq!(power_of_2_below(5), 4); + assert_eq!(power_of_2_below(7), 4); + assert_eq!(power_of_2_below(24), 16); + assert_eq!(power_of_2_below(32), 32); + assert_eq!(power_of_2_below(33), 32); + assert_eq!(power_of_2_below(63), 32); + } + + #[test] + pub fn test_power_of_2_below_large() { + let pow: u64 = 1 << 24; + for x in (pow - 20)..(pow + 20) { + assert!(power_of_2_below(x) <= x, "{}", x); + } + } +} diff --git a/eth2/fork_choice/tests/tests.rs b/eth2/fork_choice/tests/tests.rs index 7228bca10..5fb963ea5 100644 --- a/eth2/fork_choice/tests/tests.rs +++ b/eth2/fork_choice/tests/tests.rs @@ -3,7 +3,7 @@ extern crate beacon_chain; extern crate bls; extern crate db; -//extern crate env_logger; // for debugging +// extern crate env_logger; // for debugging extern crate fork_choice; extern crate hex; extern crate log; @@ -15,8 +15,10 @@ pub use beacon_chain::BeaconChain; use bls::Signature; use db::stores::{BeaconBlockStore, BeaconStateStore}; use db::MemoryDB; -//use env_logger::{Builder, Env}; -use fork_choice::{BitwiseLMDGhost, ForkChoice, ForkChoiceAlgorithm, LongestChain, SlowLMDGhost}; +// use env_logger::{Builder, Env}; +use fork_choice::{ + BitwiseLMDGhost, ForkChoice, ForkChoiceAlgorithm, LongestChain, OptimizedLMDGhost, SlowLMDGhost, +}; use ssz::ssz_encode; use std::collections::HashMap; use std::sync::Arc; @@ -27,6 +29,18 @@ use yaml_rust::yaml; // Note: We Assume the block Id's are hex-encoded. +#[test] +fn test_optimized_lmd_ghost() { + // set up logging + // Builder::from_env(Env::default().default_filter_or("trace")).init(); + + test_yaml_vectors( + ForkChoiceAlgorithm::OptimizedLMDGhost, + "tests/lmd_ghost_test_vectors.yaml", + 100, + ); +} + #[test] fn test_bitwise_lmd_ghost() { // set up logging @@ -212,6 +226,10 @@ fn setup_inital_state( // the fork choice instantiation let fork_choice: Box = match fork_choice_algo { + ForkChoiceAlgorithm::OptimizedLMDGhost => Box::new(OptimizedLMDGhost::new( + block_store.clone(), + state_store.clone(), + )), ForkChoiceAlgorithm::BitwiseLMDGhost => Box::new(BitwiseLMDGhost::new( block_store.clone(), state_store.clone(),