Fix bugs in fork choice, add more tests

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -120,7 +120,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             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<T: BeaconChainTypes> BeaconChain<T> {
         );
 
         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<T: BeaconChainTypes> BeaconChain<T> {
         Ok(self.store.get(block_root)?)
     }
 
-    /// Update the canonical head to `new_head`.
-    fn update_canonical_head(&self, new_head: CheckPoint<T::EthSpec>) -> 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<BeaconState<T::EthSpec>> {
@@ -800,16 +770,93 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                 self.metrics.fork_choice_reorg_count.inc();
             };
 
+            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,
                     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<T::EthSpec>) -> 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::<BeaconBlock>(&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<bool, Error> {

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),

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<T> = std::result::Result<T, Error>;
 
@@ -26,26 +26,40 @@ pub struct ForkChoice<T: BeaconChainTypes> {
 }
 
 impl<T: BeaconChainTypes> ForkChoice<T> {
-    pub fn new(store: Arc<T::Store>, 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<T::Store>,
+        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<T>) -> Result<Hash256> {
+        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 {
+            let (block_root, block_slot) =
-                state.current_justified_root
+                if state.current_epoch() + 1 > state.current_justified_epoch {
+                    (
+                        state.current_justified_root,
+                        start_slot(state.current_justified_epoch),
+                    )
                 } else {
-                state.finalized_root
+                    (state.finalized_root, start_slot(state.finalized_epoch))
                 };
 
             let block = chain
@@ -65,7 +79,7 @@ impl<T: BeaconChainTypes> ForkChoice<T> {
                 .get::<BeaconState<T::EthSpec>>(&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<T: BeaconChainTypes> ForkChoice<T> {
         };
 
         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<T: BeaconChainTypes> ForkChoice<T> {
             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<T: BeaconChainTypes> ForkChoice<T> {
         //  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
+        // fine becuase votes to the genesis block are not useful; all validators implicitly attest
-        // implicitly attest to genesis just by being present in the chain.
+        // to genesis just by being present in the chain.
         if block_hash != Hash256::zero() {
             let block_slot = attestation
                 .data
@@ -147,6 +161,20 @@ impl<T: BeaconChainTypes> ForkChoice<T> {
 
         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<BeaconStateError> for Error {

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<ThreadSafeReducedTree<MemoryStore, MinimalEthSpec>, 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<usize> = (0..two_thirds).collect();
+    let faulty_validators: Vec<usize> = (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"
+    );
+}

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<T> = std::result::Result<T, String>;
 
 pub trait LmdGhost<S: Store, E: EthSpec>: Send + Sync {
     /// Create a new instance, with the given `store` and `finalized_root`.
-    fn new(store: Arc<S>, finalized_root: Hash256) -> Self;
+    fn new(store: Arc<S>, 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<S: Store, E: EthSpec>: 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<F>(&self, start_block_root: Hash256, weight: F) -> Result<Hash256>
+    fn find_head<F>(
+        &self,
+        start_block_slot: Slot,
+        start_block_root: Hash256,
+        weight: F,
+    ) -> Result<Hash256>
     where
         F: Fn(usize) -> Option<u64> + Copy;
 
-    /// Provide an indication that the blockchain has been finalized at the given `finalized_root`.
+    /// Provide an indication that the blockchain has been finalized at the given `finalized_block`.
-    fn update_finalized_root(&self, finalized_root: Hash256) -> Result<()>;
+    ///
+    /// `finalized_block_root` must be the root of `finalized_block`.
+    fn update_finalized_root(
+        &self,
+        finalized_block: &BeaconBlock,
+        finalized_block_root: Hash256,
+    ) -> Result<()>;
 }

eth2/lmd_ghost/src/reduced_tree.rs

@@ -40,9 +40,9 @@ where
     T: Store,
     E: EthSpec,
 {
-    fn new(store: Arc<T>, genesis_root: Hash256) -> Self {
+    fn new(store: Arc<T>, 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<F>(&self, start_block_root: Hash256, weight_fn: F) -> SuperResult<Hash256>
+    fn find_head<F>(
+        &self,
+        start_block_slot: Slot,
+        start_block_root: Hash256,
+        weight_fn: F,
+    ) -> SuperResult<Hash256>
     where
         F: Fn(usize) -> Option<u64> + 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<T, E> {
     /// Maps validator indices to their latest votes.
     latest_votes: ElasticList<Option<Vote>>,
     /// Stores the root of the tree, used for pruning.
-    root: Hash256,
+    root: (Hash256, Slot),
     _phantom: PhantomData<E>,
 }
 
@@ -100,7 +105,7 @@ where
     T: Store,
     E: EthSpec,
 {
-    pub fn new(store: Arc<T>, genesis_root: Hash256) -> Self {
+    pub fn new(store: Arc<T>, 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,6 +171,7 @@ where
         block_hash: Hash256,
         slot: Slot,
     ) -> Result<()> {
+        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.
@@ -187,12 +199,14 @@ where
             );
 
             self.add_latest_message(validator_index, block_hash)?;
+        }
 
         Ok(())
     }
 
     pub fn update_weights_and_find_head<F>(
         &mut self,
+        start_block_slot: Slot,
         start_block_root: Hash256,
         weight_fn: F,
     ) -> Result<Hash256>
@@ -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
+                        let child = self.get_mut_node(node.children[0])?;
-                            .nodes
-                            .get_mut(&node.children[0])
-                            .ok_or_else(|| Error::MissingNode(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,7 +393,8 @@ where
         Ok(())
     }
 
-    fn add_weightless_node(&mut self, hash: Hash256) -> Result<()> {
+    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,
@@ -390,6 +407,7 @@ where
                     self.maybe_delete_node(parent_hash)?;
                 }
             }
+        }
 
         Ok(())
     }
@@ -403,11 +421,9 @@ where
             self.get_mut_node(hash)?.clone()
         };
 
-        let mut added_new_ancestor = false;
-
-        if !prev_in_tree.children.is_empty() {
         let mut added = false;
 
+        if !prev_in_tree.children.is_empty() {
             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::<BeaconState<E>>(&state_root)?
             .ok_or_else(|| Error::MissingState(state_root))
     }
+
+    fn root_slot(&self) -> Slot {
+        self.root.1
+    }
 }
 
 #[derive(Default, Clone, Debug)]