From 77fba0b98ed8a60de23ff23c70857338f559a94c Mon Sep 17 00:00:00 2001 From: Paul Hauner Date: Sun, 23 Jun 2019 14:47:23 +1000 Subject: [PATCH] Fix bugs in fork choice, add more tests --- beacon_node/beacon_chain/src/beacon_chain.rs | 125 +++++++---- beacon_node/beacon_chain/src/errors.rs | 4 + beacon_node/beacon_chain/src/fork_choice.rs | 56 +++-- beacon_node/beacon_chain/tests/tests.rs | 225 +++++++++++++++++++ eth2/lmd_ghost/src/lib.rs | 23 +- eth2/lmd_ghost/src/reduced_tree.rs | 141 +++++++----- 6 files changed, 455 insertions(+), 119 deletions(-) create mode 100644 beacon_node/beacon_chain/tests/tests.rs diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs index 2a4df66f1..0137a0746 100644 --- a/beacon_node/beacon_chain/src/beacon_chain.rs +++ b/beacon_node/beacon_chain/src/beacon_chain.rs @@ -120,7 +120,7 @@ impl BeaconChain { state: RwLock::new(genesis_state), canonical_head, genesis_block_root, - fork_choice: ForkChoice::new(store.clone(), genesis_block_root), + fork_choice: ForkChoice::new(store.clone(), &genesis_block, genesis_block_root), metrics: Metrics::new()?, store, }) @@ -145,11 +145,12 @@ impl BeaconChain { ); let last_finalized_root = p.canonical_head.beacon_state.finalized_root; + let last_finalized_block = &p.canonical_head.beacon_block; Ok(Some(BeaconChain { spec, slot_clock, - fork_choice: ForkChoice::new(store.clone(), last_finalized_root), + fork_choice: ForkChoice::new(store.clone(), last_finalized_block, last_finalized_root), op_pool: OperationPool::default(), canonical_head: RwLock::new(p.canonical_head), state: RwLock::new(p.state), @@ -239,37 +240,6 @@ impl BeaconChain { Ok(self.store.get(block_root)?) } - /// Update the canonical head to `new_head`. - fn update_canonical_head(&self, new_head: CheckPoint) -> Result<(), Error> { - // Update the checkpoint that stores the head of the chain at the time it received the - // block. - *self.canonical_head.write() = new_head; - - // Update the always-at-the-present-slot state we keep around for performance gains. - *self.state.write() = { - let mut state = self.canonical_head.read().beacon_state.clone(); - - let present_slot = match self.slot_clock.present_slot() { - Ok(Some(slot)) => slot, - _ => return Err(Error::UnableToReadSlot), - }; - - // If required, transition the new state to the present slot. - for _ in state.slot.as_u64()..present_slot.as_u64() { - per_slot_processing(&mut state, &self.spec)?; - } - - state.build_all_caches(&self.spec)?; - - state - }; - - // Save `self` to `self.store`. - self.persist()?; - - Ok(()) - } - /// Returns a read-lock guarded `BeaconState` which is the `canonical_head` that has been /// updated to match the current slot clock. pub fn current_state(&self) -> RwLockReadGuard> { @@ -800,17 +770,94 @@ impl BeaconChain { self.metrics.fork_choice_reorg_count.inc(); }; - self.update_canonical_head(CheckPoint { - beacon_block, - beacon_block_root, - beacon_state, - beacon_state_root, - })?; + let old_finalized_epoch = self.head().beacon_state.finalized_epoch; + let new_finalized_epoch = beacon_state.finalized_epoch; + let finalized_root = beacon_state.finalized_root; + + // Never revert back past a finalized epoch. + if new_finalized_epoch < old_finalized_epoch { + Err(Error::RevertedFinalizedEpoch { + previous_epoch: old_finalized_epoch, + new_epoch: new_finalized_epoch, + }) + } else { + self.update_canonical_head(CheckPoint { + beacon_block: beacon_block, + beacon_block_root, + beacon_state, + beacon_state_root, + })?; + + if new_finalized_epoch != old_finalized_epoch { + self.after_finalization(old_finalized_epoch, finalized_root)?; + } + + Ok(()) + } + } else { + Ok(()) } + } + + /// Update the canonical head to `new_head`. + fn update_canonical_head(&self, new_head: CheckPoint) -> Result<(), Error> { + // Update the checkpoint that stores the head of the chain at the time it received the + // block. + *self.canonical_head.write() = new_head; + + // Update the always-at-the-present-slot state we keep around for performance gains. + *self.state.write() = { + let mut state = self.canonical_head.read().beacon_state.clone(); + + let present_slot = match self.slot_clock.present_slot() { + Ok(Some(slot)) => slot, + _ => return Err(Error::UnableToReadSlot), + }; + + // If required, transition the new state to the present slot. + for _ in state.slot.as_u64()..present_slot.as_u64() { + per_slot_processing(&mut state, &self.spec)?; + } + + state.build_all_caches(&self.spec)?; + + state + }; + + // Save `self` to `self.store`. + self.persist()?; Ok(()) } + /// Called after `self` has had a new block finalized. + /// + /// Performs pruning and finality-based optimizations. + fn after_finalization( + &self, + old_finalized_epoch: Epoch, + finalized_block_root: Hash256, + ) -> Result<(), Error> { + let finalized_block = self + .store + .get::(&finalized_block_root)? + .ok_or_else(|| Error::MissingBeaconBlock(finalized_block_root))?; + + let new_finalized_epoch = finalized_block.slot.epoch(T::EthSpec::slots_per_epoch()); + + if new_finalized_epoch < old_finalized_epoch { + Err(Error::RevertedFinalizedEpoch { + previous_epoch: old_finalized_epoch, + new_epoch: new_finalized_epoch, + }) + } else { + self.fork_choice + .process_finalization(&finalized_block, finalized_block_root)?; + + Ok(()) + } + } + /// Returns `true` if the given block root has not been processed. pub fn is_new_block_root(&self, beacon_block_root: &Hash256) -> Result { Ok(!self.store.exists::(beacon_block_root)?) diff --git a/beacon_node/beacon_chain/src/errors.rs b/beacon_node/beacon_chain/src/errors.rs index 157d774c6..0d619d7f2 100644 --- a/beacon_node/beacon_chain/src/errors.rs +++ b/beacon_node/beacon_chain/src/errors.rs @@ -19,6 +19,10 @@ pub enum BeaconChainError { InsufficientValidators, BadRecentBlockRoots, UnableToReadSlot, + RevertedFinalizedEpoch { + previous_epoch: Epoch, + new_epoch: Epoch, + }, BeaconStateError(BeaconStateError), DBInconsistent(String), DBError(store::Error), diff --git a/beacon_node/beacon_chain/src/fork_choice.rs b/beacon_node/beacon_chain/src/fork_choice.rs index 7aba6fdf5..389e5b46b 100644 --- a/beacon_node/beacon_chain/src/fork_choice.rs +++ b/beacon_node/beacon_chain/src/fork_choice.rs @@ -3,7 +3,7 @@ use lmd_ghost::LmdGhost; use state_processing::common::get_attesting_indices_unsorted; use std::sync::Arc; use store::{Error as StoreError, Store}; -use types::{Attestation, BeaconBlock, BeaconState, BeaconStateError, EthSpec, Hash256}; +use types::{Attestation, BeaconBlock, BeaconState, BeaconStateError, Epoch, EthSpec, Hash256}; type Result = std::result::Result; @@ -26,27 +26,41 @@ pub struct ForkChoice { } impl ForkChoice { - pub fn new(store: Arc, genesis_block_root: Hash256) -> Self { + /// Instantiate a new fork chooser. + /// + /// "Genesis" does not necessarily need to be the absolute genesis, it can be some finalized + /// block. + pub fn new( + store: Arc, + genesis_block: &BeaconBlock, + genesis_block_root: Hash256, + ) -> Self { Self { - backend: T::LmdGhost::new(store, genesis_block_root), + backend: T::LmdGhost::new(store, genesis_block, genesis_block_root), genesis_block_root, } } pub fn find_head(&self, chain: &BeaconChain) -> Result { + let start_slot = |epoch: Epoch| epoch.start_slot(T::EthSpec::slots_per_epoch()); + // 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 (start_state, start_block_root, start_block_slot) = { 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_root, block_slot) = + if state.current_epoch() + 1 > state.current_justified_epoch { + ( + state.current_justified_root, + start_slot(state.current_justified_epoch), + ) + } else { + (state.finalized_root, start_slot(state.finalized_epoch)) + }; let block = chain .store @@ -65,7 +79,7 @@ impl ForkChoice { .get::>(&block.state_root)? .ok_or_else(|| Error::MissingState(block.state_root))?; - (state, block_root) + (state, block_root, block_slot) }; // A function that returns the weight for some validator index. @@ -77,7 +91,7 @@ impl ForkChoice { }; self.backend - .find_head(start_block_root, weight) + .find_head(start_block_slot, start_block_root, weight) .map_err(Into::into) } @@ -101,7 +115,7 @@ impl ForkChoice { self.process_attestation_from_block(state, attestation)?; } - self.backend.process_block(block_root, block.slot)?; + self.backend.process_block(block, block_root)?; Ok(()) } @@ -131,8 +145,8 @@ impl ForkChoice { // 2. Ignore all attestations to the zero hash. // // (1) becomes weird once we hit finality and fork choice drops the genesis block. (2) is - // fine becuase votes to the genesis block are not usefully, all validators already - // implicitly attest to genesis just by being present in the chain. + // fine becuase votes to the genesis block are not useful; all validators implicitly attest + // to genesis just by being present in the chain. if block_hash != Hash256::zero() { let block_slot = attestation .data @@ -147,6 +161,20 @@ impl ForkChoice { Ok(()) } + + /// Inform the fork choice that the given block (and corresponding root) have been finalized so + /// it may prune it's storage. + /// + /// `finalized_block_root` must be the root of `finalized_block`. + pub fn process_finalization( + &self, + finalized_block: &BeaconBlock, + finalized_block_root: Hash256, + ) -> Result<()> { + self.backend + .update_finalized_root(finalized_block, finalized_block_root) + .map_err(Into::into) + } } impl From for Error { diff --git a/beacon_node/beacon_chain/tests/tests.rs b/beacon_node/beacon_chain/tests/tests.rs new file mode 100644 index 000000000..5181f1e76 --- /dev/null +++ b/beacon_node/beacon_chain/tests/tests.rs @@ -0,0 +1,225 @@ +use beacon_chain::test_utils::{AttestationStrategy, BeaconChainHarness, BlockStrategy}; +use lmd_ghost::ThreadSafeReducedTree; +use store::MemoryStore; +use types::{EthSpec, MinimalEthSpec, Slot}; + +// Should ideally be divisible by 3. +pub const VALIDATOR_COUNT: usize = 24; + +fn get_harness( + validator_count: usize, +) -> BeaconChainHarness, MinimalEthSpec> { + let harness = BeaconChainHarness::new(validator_count); + + // Move past the zero slot. + harness.advance_slot(); + + harness +} + +#[test] +fn fork() { + let harness = get_harness(VALIDATOR_COUNT); + + let two_thirds = (VALIDATOR_COUNT / 3) * 2; + let delay = MinimalEthSpec::default_spec().min_attestation_inclusion_delay as usize; + + let honest_validators: Vec = (0..two_thirds).collect(); + let faulty_validators: Vec = (two_thirds..VALIDATOR_COUNT).collect(); + + let initial_blocks = delay + 1; + let honest_fork_blocks = delay + 1; + let faulty_fork_blocks = delay + 2; + + // Build an initial chain were all validators agree. + harness.extend_chain( + initial_blocks, + BlockStrategy::OnCanonicalHead, + AttestationStrategy::AllValidators, + ); + + // Move to the next slot so we may produce some more blocks on the head. + harness.advance_slot(); + + // Extend the chain with blocks where only honest validators agree. + let honest_head = harness.extend_chain( + honest_fork_blocks, + BlockStrategy::OnCanonicalHead, + AttestationStrategy::SomeValidators(honest_validators.clone()), + ); + + // Go back to the last block where all agreed, and build blocks upon it where only faulty nodes + // agree. + let faulty_head = harness.extend_chain( + faulty_fork_blocks, + BlockStrategy::ForkCanonicalChainAt { + previous_slot: Slot::from(initial_blocks), + first_slot: Slot::from(initial_blocks + 2), + }, + AttestationStrategy::SomeValidators(faulty_validators.clone()), + ); + + assert!(honest_head != faulty_head, "forks should be distinct"); + + let state = &harness.chain.head().beacon_state; + + assert_eq!( + state.slot, + Slot::from(initial_blocks + honest_fork_blocks), + "head should be at the current slot" + ); + + assert_eq!( + harness.chain.head().beacon_block_root, + honest_head, + "the honest chain should be the canonical chain" + ); +} + +#[test] +fn finalizes_with_full_participation() { + let num_blocks_produced = MinimalEthSpec::slots_per_epoch() * 5; + + let harness = get_harness(VALIDATOR_COUNT); + + harness.extend_chain( + num_blocks_produced as usize, + BlockStrategy::OnCanonicalHead, + AttestationStrategy::AllValidators, + ); + + let state = &harness.chain.head().beacon_state; + + assert_eq!( + state.slot, num_blocks_produced, + "head should be at the current slot" + ); + assert_eq!( + state.current_epoch(), + num_blocks_produced / MinimalEthSpec::slots_per_epoch(), + "head should be at the expected epoch" + ); + assert_eq!( + state.current_justified_epoch, + state.current_epoch() - 1, + "the head should be justified one behind the current epoch" + ); + assert_eq!( + state.finalized_epoch, + state.current_epoch() - 2, + "the head should be finalized two behind the current epoch" + ); +} + +#[test] +fn finalizes_with_two_thirds_participation() { + let num_blocks_produced = MinimalEthSpec::slots_per_epoch() * 5; + + let harness = get_harness(VALIDATOR_COUNT); + + let two_thirds = (VALIDATOR_COUNT / 3) * 2; + let attesters = (0..two_thirds).collect(); + + harness.extend_chain( + num_blocks_produced as usize, + BlockStrategy::OnCanonicalHead, + AttestationStrategy::SomeValidators(attesters), + ); + + let state = &harness.chain.head().beacon_state; + + assert_eq!( + state.slot, num_blocks_produced, + "head should be at the current slot" + ); + assert_eq!( + state.current_epoch(), + num_blocks_produced / MinimalEthSpec::slots_per_epoch(), + "head should be at the expected epoch" + ); + + // Note: the 2/3rds tests are not justifying the immediately prior epochs because the + // `MIN_ATTESTATION_INCLUSION_DELAY` is preventing an adequate number of attestations being + // included in blocks during that epoch. + + assert_eq!( + state.current_justified_epoch, + state.current_epoch() - 2, + "the head should be justified two behind the current epoch" + ); + assert_eq!( + state.finalized_epoch, + state.current_epoch() - 4, + "the head should be finalized three behind the current epoch" + ); +} + +#[test] +fn does_not_finalize_with_less_than_two_thirds_participation() { + let num_blocks_produced = MinimalEthSpec::slots_per_epoch() * 5; + + let harness = get_harness(VALIDATOR_COUNT); + + let two_thirds = (VALIDATOR_COUNT / 3) * 2; + let less_than_two_thirds = two_thirds - 1; + let attesters = (0..less_than_two_thirds).collect(); + + harness.extend_chain( + num_blocks_produced as usize, + BlockStrategy::OnCanonicalHead, + AttestationStrategy::SomeValidators(attesters), + ); + + let state = &harness.chain.head().beacon_state; + + assert_eq!( + state.slot, num_blocks_produced, + "head should be at the current slot" + ); + assert_eq!( + state.current_epoch(), + num_blocks_produced / MinimalEthSpec::slots_per_epoch(), + "head should be at the expected epoch" + ); + assert_eq!( + state.current_justified_epoch, 0, + "no epoch should have been justified" + ); + assert_eq!( + state.finalized_epoch, 0, + "no epoch should have been finalized" + ); +} + +#[test] +fn does_not_finalize_without_attestation() { + let num_blocks_produced = MinimalEthSpec::slots_per_epoch() * 5; + + let harness = get_harness(VALIDATOR_COUNT); + + harness.extend_chain( + num_blocks_produced as usize, + BlockStrategy::OnCanonicalHead, + AttestationStrategy::SomeValidators(vec![]), + ); + + let state = &harness.chain.head().beacon_state; + + assert_eq!( + state.slot, num_blocks_produced, + "head should be at the current slot" + ); + assert_eq!( + state.current_epoch(), + num_blocks_produced / MinimalEthSpec::slots_per_epoch(), + "head should be at the expected epoch" + ); + assert_eq!( + state.current_justified_epoch, 0, + "no epoch should have been justified" + ); + assert_eq!( + state.finalized_epoch, 0, + "no epoch should have been finalized" + ); +} diff --git a/eth2/lmd_ghost/src/lib.rs b/eth2/lmd_ghost/src/lib.rs index 8f5d0b529..dd413e2eb 100644 --- a/eth2/lmd_ghost/src/lib.rs +++ b/eth2/lmd_ghost/src/lib.rs @@ -2,7 +2,7 @@ mod reduced_tree; use std::sync::Arc; use store::Store; -use types::{EthSpec, Hash256, Slot}; +use types::{BeaconBlock, EthSpec, Hash256, Slot}; pub use reduced_tree::ThreadSafeReducedTree; @@ -10,7 +10,7 @@ pub type Result = std::result::Result; pub trait LmdGhost: Send + Sync { /// Create a new instance, with the given `store` and `finalized_root`. - fn new(store: Arc, finalized_root: Hash256) -> Self; + fn new(store: Arc, finalized_block: &BeaconBlock, finalized_root: Hash256) -> Self; /// Process an attestation message from some validator that attests to some `block_hash` /// representing a block at some `block_slot`. @@ -22,14 +22,25 @@ pub trait LmdGhost: Send + Sync { ) -> Result<()>; /// Process a block that was seen on the network. - fn process_block(&self, block_hash: Hash256, block_slot: Slot) -> Result<()>; + fn process_block(&self, block: &BeaconBlock, block_hash: Hash256) -> Result<()>; /// Returns the head of the chain, starting the search at `start_block_root` and moving upwards /// (in block height). - fn find_head(&self, start_block_root: Hash256, weight: F) -> Result + fn find_head( + &self, + start_block_slot: Slot, + start_block_root: Hash256, + weight: F, + ) -> Result where F: Fn(usize) -> Option + Copy; - /// Provide an indication that the blockchain has been finalized at the given `finalized_root`. - fn update_finalized_root(&self, finalized_root: Hash256) -> Result<()>; + /// Provide an indication that the blockchain has been finalized at the given `finalized_block`. + /// + /// `finalized_block_root` must be the root of `finalized_block`. + fn update_finalized_root( + &self, + finalized_block: &BeaconBlock, + finalized_block_root: Hash256, + ) -> Result<()>; } diff --git a/eth2/lmd_ghost/src/reduced_tree.rs b/eth2/lmd_ghost/src/reduced_tree.rs index 2d70ea145..82547bbf6 100644 --- a/eth2/lmd_ghost/src/reduced_tree.rs +++ b/eth2/lmd_ghost/src/reduced_tree.rs @@ -40,9 +40,9 @@ where T: Store, E: EthSpec, { - fn new(store: Arc, genesis_root: Hash256) -> Self { + fn new(store: Arc, genesis_block: &BeaconBlock, genesis_root: Hash256) -> Self { ThreadSafeReducedTree { - core: RwLock::new(ReducedTree::new(store, genesis_root)), + core: RwLock::new(ReducedTree::new(store, genesis_block, genesis_root)), } } @@ -59,27 +59,32 @@ where } /// Process a block that was seen on the network. - fn process_block(&self, block_hash: Hash256, _block_slot: Slot) -> SuperResult<()> { + fn process_block(&self, block: &BeaconBlock, block_hash: Hash256) -> SuperResult<()> { self.core .write() - .add_weightless_node(block_hash) + .add_weightless_node(block.slot, block_hash) .map_err(|e| format!("process_block failed: {:?}", e)) } - fn find_head(&self, start_block_root: Hash256, weight_fn: F) -> SuperResult + fn find_head( + &self, + start_block_slot: Slot, + start_block_root: Hash256, + weight_fn: F, + ) -> SuperResult where F: Fn(usize) -> Option + Copy, { self.core .write() - .update_weights_and_find_head(start_block_root, weight_fn) + .update_weights_and_find_head(start_block_slot, start_block_root, weight_fn) .map_err(|e| format!("find_head failed: {:?}", e)) } - fn update_finalized_root(&self, new_root: Hash256) -> SuperResult<()> { + fn update_finalized_root(&self, new_block: &BeaconBlock, new_root: Hash256) -> SuperResult<()> { self.core .write() - .update_root(new_root) + .update_root(new_block.slot, new_root) .map_err(|e| format!("update_finalized_root failed: {:?}", e)) } } @@ -91,7 +96,7 @@ struct ReducedTree { /// Maps validator indices to their latest votes. latest_votes: ElasticList>, /// Stores the root of the tree, used for pruning. - root: Hash256, + root: (Hash256, Slot), _phantom: PhantomData, } @@ -100,7 +105,7 @@ where T: Store, E: EthSpec, { - pub fn new(store: Arc, genesis_root: Hash256) -> Self { + pub fn new(store: Arc, genesis_block: &BeaconBlock, genesis_root: Hash256) -> Self { let mut nodes = HashMap::new(); // Insert the genesis node. @@ -116,12 +121,12 @@ where store, nodes, latest_votes: ElasticList::default(), - root: genesis_root, + root: (genesis_root, genesis_block.slot), _phantom: PhantomData, } } - pub fn update_root(&mut self, new_root: Hash256) -> Result<()> { + pub fn update_root(&mut self, new_slot: Slot, new_root: Hash256) -> Result<()> { if !self.nodes.contains_key(&new_root) { let node = Node { block_hash: new_root, @@ -132,16 +137,22 @@ where self.add_node(node)?; } - self.retain_subtree(self.root, new_root)?; + self.retain_subtree(self.root.0, new_root)?; - self.root = new_root; + self.root = (new_root, new_slot); + + let root_node = self.get_mut_node(new_root)?; + root_node.parent_hash = None; Ok(()) } + /// Removes `current_hash` and all decendants, except `subtree_hash` and all nodes + /// which have `subtree_hash` as an ancestor. + /// + /// In effect, prunes the tree so that only decendants of `subtree_hash` exist. fn retain_subtree(&mut self, current_hash: Hash256, subtree_hash: Hash256) -> Result<()> { if current_hash != subtree_hash { - // Clone satisifies the borrow checker. let children = self.get_node(current_hash)?.children.clone(); for child_hash in children { @@ -160,39 +171,42 @@ where block_hash: Hash256, slot: Slot, ) -> Result<()> { - if let Some(previous_vote) = self.latest_votes.get(validator_index) { - if previous_vote.slot > slot { - // Given vote is earier than known vote, nothing to do. - return Ok(()); - } else if previous_vote.slot == slot && previous_vote.hash == block_hash { - // Given vote is identical to known vote, nothing to do. - return Ok(()); - } else if previous_vote.slot == slot && previous_vote.hash != block_hash { - // Vote is an equivocation (double-vote), ignore it. - // - // TODO: this is slashable. - return Ok(()); - } else { - // Given vote is newer or different to current vote, replace the current vote. - self.remove_latest_message(validator_index)?; + if slot >= self.root_slot() { + if let Some(previous_vote) = self.latest_votes.get(validator_index) { + if previous_vote.slot > slot { + // Given vote is earier than known vote, nothing to do. + return Ok(()); + } else if previous_vote.slot == slot && previous_vote.hash == block_hash { + // Given vote is identical to known vote, nothing to do. + return Ok(()); + } else if previous_vote.slot == slot && previous_vote.hash != block_hash { + // Vote is an equivocation (double-vote), ignore it. + // + // TODO: this is slashable. + return Ok(()); + } else { + // Given vote is newer or different to current vote, replace the current vote. + self.remove_latest_message(validator_index)?; + } } + + self.latest_votes.insert( + validator_index, + Some(Vote { + slot, + hash: block_hash, + }), + ); + + self.add_latest_message(validator_index, block_hash)?; } - self.latest_votes.insert( - validator_index, - Some(Vote { - slot, - hash: block_hash, - }), - ); - - self.add_latest_message(validator_index, block_hash)?; - Ok(()) } pub fn update_weights_and_find_head( &mut self, + start_block_slot: Slot, start_block_root: Hash256, weight_fn: F, ) -> Result @@ -203,7 +217,7 @@ where // // In this case, we add a weightless node at `start_block_root`. if !self.nodes.contains_key(&start_block_root) { - self.add_weightless_node(start_block_root)?; + self.add_weightless_node(start_block_slot, start_block_root)?; }; let _root_weight = self.update_weight(start_block_root, weight_fn)?; @@ -289,13 +303,15 @@ where // // Load the child of the node and set it's parent to be the parent of this // node (viz., graft the node's child to the node's parent) - let child = self - .nodes - .get_mut(&node.children[0]) - .ok_or_else(|| Error::MissingNode(node.children[0]))?; - + let child = self.get_mut_node(node.children[0])?; child.parent_hash = node.parent_hash; + // Graft the parent of this node to it's child. + if let Some(parent_hash) = node.parent_hash { + let parent = self.get_mut_node(parent_hash)?; + parent.replace_child(node.block_hash, node.children[0])?; + } + true } else if node.children.len() == 0 { // A node which has no children may be deleted and potentially it's parent @@ -377,17 +393,19 @@ where Ok(()) } - fn add_weightless_node(&mut self, hash: Hash256) -> Result<()> { - if !self.nodes.contains_key(&hash) { - let node = Node { - block_hash: hash, - ..Node::default() - }; + fn add_weightless_node(&mut self, slot: Slot, hash: Hash256) -> Result<()> { + if slot >= self.root_slot() { + if !self.nodes.contains_key(&hash) { + let node = Node { + block_hash: hash, + ..Node::default() + }; - self.add_node(node)?; + self.add_node(node)?; - if let Some(parent_hash) = self.get_node(hash)?.parent_hash { - self.maybe_delete_node(parent_hash)?; + if let Some(parent_hash) = self.get_node(hash)?.parent_hash { + self.maybe_delete_node(parent_hash)?; + } } } @@ -403,11 +421,9 @@ where self.get_mut_node(hash)?.clone() }; - let mut added_new_ancestor = false; + let mut added = false; if !prev_in_tree.children.is_empty() { - let mut added = false; - for &child_hash in &prev_in_tree.children { if self .iter_ancestors(child_hash)? @@ -418,6 +434,7 @@ where child.parent_hash = Some(node.block_hash); node.children.push(child_hash); prev_in_tree.replace_child(child_hash, node.block_hash)?; + node.parent_hash = Some(prev_in_tree.block_hash); added = true; @@ -446,7 +463,7 @@ where self.nodes .insert(common_ancestor.block_hash, common_ancestor); - added_new_ancestor = true; + added = true; break; } @@ -454,7 +471,7 @@ where } } - if !added_new_ancestor { + if !added { node.parent_hash = Some(prev_in_tree.block_hash); prev_in_tree.children.push(node.block_hash); } @@ -558,6 +575,10 @@ where .get::>(&state_root)? .ok_or_else(|| Error::MissingState(state_root)) } + + fn root_slot(&self) -> Slot { + self.root.1 + } } #[derive(Default, Clone, Debug)]