Add first pass at fixing filter

eth2/proto_array_fork_choice/src/lib.rs

@@ -30,12 +30,13 @@ pub enum Error {
         indices: usize,
     },
     RevertedFinalizedEpoch,
-    InvalidFindHeadStartRoot {
+    InvalidBestNode {
         justified_epoch: Epoch,
         finalized_epoch: Epoch,
         node_justified_epoch: Epoch,
         node_finalized_epoch: Epoch,
     },
+    BestDescendantWithoutBestChild,
 }
 
 #[derive(Default, PartialEq, Clone, Encode, Decode)]
@@ -102,7 +103,6 @@ impl ProtoArrayForkChoice {
             ffg_update_required: false,
             justified_epoch,
             finalized_epoch,
-            finalized_root,
             nodes: Vec::with_capacity(1),
             indices: HashMap::with_capacity(1),
         };

eth2/proto_array_fork_choice/src/proto_array.rs

@@ -37,7 +37,6 @@ pub struct ProtoArray {
     pub ffg_update_required: bool,
     pub justified_epoch: Epoch,
     pub finalized_epoch: Epoch,
-    pub finalized_root: Hash256,
     pub nodes: Vec<ProtoNode>,
     pub indices: HashMap<Hash256, usize>,
 }
@@ -131,6 +130,9 @@ impl ProtoArray {
                 // Back-propogate the nodes delta to its parent.
                 *parent_delta += node_delta;
 
+                self.maybe_update_best_child_and_descendant(parent_index, node_index)?;
+
+                /*
                 let is_viable_for_head = self
                     .nodes
                     .get(node_index)
@@ -206,6 +208,7 @@ impl ProtoArray {
                     // this code only runs if the current node is viable for the head).
                     self.set_best_child(parent_index, node_index)?;
                 };
+                */
             }
         }
 
@@ -220,7 +223,7 @@ impl ProtoArray {
     pub fn on_new_block(
         &mut self,
         root: Hash256,
-        parent: Option<Hash256>,
+        parent_opt: Option<Hash256>,
         justified_epoch: Epoch,
         finalized_epoch: Epoch,
     ) -> Result<(), Error> {
@@ -228,7 +231,7 @@ impl ProtoArray {
 
         let node = ProtoNode {
             root,
-            parent: parent.and_then(|parent_root| self.indices.get(&parent_root).copied()),
+            parent: parent_opt.and_then(|parent| self.indices.get(&parent).copied()),
             justified_epoch,
             finalized_epoch,
             weight: 0,
@@ -239,6 +242,11 @@ impl ProtoArray {
         self.indices.insert(node.root, node_index);
         self.nodes.push(node.clone());
 
+        if let Some(parent_index) = node.parent {
+            self.maybe_update_best_child_and_descendant(parent_index, node_index)?;
+        }
+
+        /*
         // If the blocks justified and finalized epochs match our values, then try and see if it
         // becomes the best child.
         if self.node_is_viable_for_head(&node) {
@@ -262,6 +270,7 @@ impl ProtoArray {
                 };
             }
         }
+        */
 
         Ok(())
     }
@@ -279,24 +288,13 @@ impl ProtoArray {
             .indices
             .get(justified_root)
             .copied()
-            .ok_or_else(|| Error::JustifiedNodeUnknown(self.finalized_root))?;
+            .ok_or_else(|| Error::JustifiedNodeUnknown(*justified_root))?;
 
         let justified_node = self
             .nodes
             .get(justified_index)
             .ok_or_else(|| Error::InvalidJustifiedIndex(justified_index))?;
 
-        // It is a logic error to try and find the head starting from a block that does not match
-        // the filter.
-        if !self.node_is_viable_for_head(&justified_node) {
-            return Err(Error::InvalidFindHeadStartRoot {
-                justified_epoch: self.justified_epoch,
-                finalized_epoch: self.finalized_epoch,
-                node_justified_epoch: justified_node.justified_epoch,
-                node_finalized_epoch: justified_node.finalized_epoch,
-            });
-        }
-
         let best_descendant_index = justified_node
             .best_descendant
             .unwrap_or_else(|| justified_index);
@@ -306,6 +304,19 @@ impl ProtoArray {
             .get(best_descendant_index)
             .ok_or_else(|| Error::InvalidBestDescendant(best_descendant_index))?;
 
+        // It is a logic error to try and find the head starting from a block that does not match
+        // the filter.
+        if !self.node_is_viable_for_head(&best_node) {
+            return Err(Error::InvalidBestNode {
+                justified_epoch: self.justified_epoch,
+                finalized_epoch: self.finalized_epoch,
+                node_justified_epoch: justified_node.justified_epoch,
+                node_finalized_epoch: justified_node.finalized_epoch,
+            });
+        }
+
+        // dbg!(&self.nodes);
+
         Ok(best_node.root)
     }
 
@@ -327,24 +338,19 @@ impl ProtoArray {
         finalized_epoch: Epoch,
         finalized_root: Hash256,
     ) -> Result<(), Error> {
-        if finalized_epoch == self.finalized_epoch && self.finalized_root != finalized_root {
+        if finalized_epoch < self.finalized_epoch {
-            // It's illegal to swap finalized roots on the same epoch (this is reverting a
-            // finalized block).
-            return Err(Error::InvalidFinalizedRootChange);
-        } else if finalized_epoch < self.finalized_epoch {
             // It's illegal to swap to an earlier finalized root (this is assumed to be reverting a
             // finalized block).
             return Err(Error::RevertedFinalizedEpoch);
         } else if finalized_epoch != self.finalized_epoch {
             self.finalized_epoch = finalized_epoch;
-            self.finalized_root = finalized_root;
             self.ffg_update_required = true;
         }
 
         let finalized_index = *self
             .indices
-            .get(&self.finalized_root)
+            .get(&finalized_root)
-            .ok_or_else(|| Error::FinalizedNodeUnknown(self.finalized_root))?;
+            .ok_or_else(|| Error::FinalizedNodeUnknown(finalized_root))?;
 
         if finalized_index < self.prune_threshold {
             // Pruning at small numbers incurs more cost than benefit.
@@ -397,6 +403,189 @@ impl ProtoArray {
         Ok(())
     }
 
+    fn maybe_update_best_child_and_descendant(
+        &mut self,
+        parent_index: usize,
+        child_index: usize,
+    ) -> Result<(), Error> {
+        let child = self
+            .nodes
+            .get(child_index)
+            .ok_or_else(|| Error::InvalidNodeIndex(child_index))?;
+
+        let parent = self
+            .nodes
+            .get(parent_index)
+            .ok_or_else(|| Error::InvalidNodeIndex(parent_index))?;
+
+        let child_leads_to_viable_head = self.node_leads_to_viable_head(&child)?;
+
+        let change_to_none = (None, None);
+        let change_to_child = (
+            Some(child_index),
+            child.best_descendant.or(Some(child_index)),
+        );
+        let no_change = (parent.best_child, parent.best_descendant);
+
+        let (new_best_child, new_best_descendant) =
+            match (parent.best_child, parent.best_descendant) {
+                (None, None) => {
+                    if child_leads_to_viable_head {
+                        change_to_child
+                    } else {
+                        no_change
+                    }
+                }
+                (Some(best_child_index), Some(best_descendant_index)) => {
+                    if best_child_index == child_index && !child_leads_to_viable_head {
+                        change_to_none
+                    } else if best_child_index == child_index {
+                        change_to_child
+                    } else {
+                        let best_child = self
+                            .nodes
+                            .get(best_child_index)
+                            .ok_or_else(|| Error::InvalidBestDescendant(best_child_index))?;
+                        /*
+                        let best_descendant = self
+                            .nodes
+                            .get(best_descendant_index)
+                            .ok_or_else(|| Error::InvalidBestDescendant(best_descendant_index))?;
+                        */
+
+                        let child_leads_to_viable_head = self.node_leads_to_viable_head(&child)?;
+                        let best_child_leads_to_viable_head =
+                            self.node_leads_to_viable_head(&best_child)?;
+
+                        if child_leads_to_viable_head && !best_child_leads_to_viable_head {
+                            change_to_child
+                        } else if !child_leads_to_viable_head && best_child_leads_to_viable_head {
+                            no_change
+                        } else if child.weight == best_child.weight {
+                            if child.root >= best_child.root {
+                                change_to_child
+                            } else {
+                                no_change
+                            }
+                        } else {
+                            if child.weight >= best_child.weight {
+                                change_to_child
+                            } else {
+                                no_change
+                            }
+                        }
+                    }
+                }
+                _ => return Err(Error::BestDescendantWithoutBestChild),
+            };
+
+        /*
+        dbg!((
+            child_index,
+            parent_index,
+            new_best_child,
+            new_best_descendant
+        ));
+        */
+
+        let parent = self
+            .nodes
+            .get_mut(parent_index)
+            .ok_or_else(|| Error::InvalidNodeIndex(parent_index))?;
+
+        parent.best_child = new_best_child;
+        parent.best_descendant = new_best_descendant;
+
+        /*
+        let new_best_descendant = if let Some(parent_best_descendant_index) = parent.best_descendant
+        {
+            if parent_best_descendant_index == child_index && !child_leads_to_viable_head {
+                None
+            } else if parent_best_descendant_index != child_index {
+                let parent_best_descendant = self
+                    .nodes
+                    .get(parent_best_descendant_index)
+                    .ok_or_else(|| Error::InvalidBestDescendant(parent_best_descendant_index))?;
+
+                let child_leads_to_viable_head = self.node_leads_to_viable_head(&child)?;
+                let parent_best_descendant_leads_to_viable_head =
+                    self.node_leads_to_viable_head(&parent_best_descendant)?;
+
+                if child_leads_to_viable_head && !parent_best_descendant_leads_to_viable_head {
+                    Some(child_index)
+                } else if !child_leads_to_viable_head && parent_best_descendant_leads_to_viable_head
+                {
+                    Some(parent_best_descendant_index)
+                } else if child.weight == parent_best_descendant.weight {
+                    if child.root >= parent_best_descendant.root {
+                        Some(child_index)
+                    } else {
+                        Some(parent_best_descendant_index)
+                    }
+                } else {
+                    if child.weight >= parent_best_descendant.weight {
+                        Some(child_index)
+                    } else {
+                        Some(parent_best_descendant_index)
+                    }
+                }
+            } else {
+                Some(child_index)
+            }
+        } else {
+            if child_leads_to_viable_head {
+                // If the parent does not have a best-descendant and the child is viable, then it's
+                // the best by default.
+                Some(child_index)
+            } else {
+                // If the parent does not have a best-descendant but the child is not viable, then
+                // leave the parent with no best-descendant.
+                None
+            }
+        };
+
+        let parent = self
+            .nodes
+            .get_mut(parent_index)
+            .ok_or_else(|| Error::InvalidNodeIndex(parent_index))?;
+
+        dbg!(new_best_descendant);
+
+        match new_best_descendant {
+            None => {
+                parent.best_child = None;
+                parent.best_descendant = None;
+            }
+            Some(index) if index == child_index => {
+                parent.best_child = Some(child_index);
+                parent.best_descendant = match child_best_descendant {
+                    Some(index) => Some(index),
+                    None => Some(child_index),
+                };
+            }
+            _ => {}
+        }
+        */
+
+        Ok(())
+    }
+
+    fn node_leads_to_viable_head(&self, node: &ProtoNode) -> Result<bool, Error> {
+        let best_descendant_is_viable_for_head =
+            if let Some(best_descendant_index) = node.best_descendant {
+                let best_descendant = self
+                    .nodes
+                    .get(best_descendant_index)
+                    .ok_or_else(|| Error::InvalidBestDescendant(best_descendant_index))?;
+
+                self.node_is_viable_for_head(best_descendant)
+            } else {
+                false
+            };
+
+        Ok(best_descendant_is_viable_for_head || self.node_is_viable_for_head(node))
+    }
+
     /// Sets the node at `parent_index` to have a best-child pointing to `child_index`. Also
     /// updates the best-descendant.
     fn set_best_child(&mut self, parent_index: usize, child_index: usize) -> Result<(), Error> {

eth2/proto_array_fork_choice/src/ssz_container.rs

@@ -16,7 +16,6 @@ pub struct SszContainer {
     ffg_update_required: bool,
     justified_epoch: Epoch,
     finalized_epoch: Epoch,
-    finalized_root: Hash256,
     nodes: Vec<ProtoNode>,
     indices: Vec<(Hash256, usize)>,
 }
@@ -32,7 +31,6 @@ impl From<&ProtoArrayForkChoice> for SszContainer {
             ffg_update_required: proto_array.ffg_update_required,
             justified_epoch: proto_array.justified_epoch,
             finalized_epoch: proto_array.finalized_epoch,
-            finalized_root: proto_array.finalized_root,
             nodes: proto_array.nodes.clone(),
             indices: proto_array.indices.iter().map(|(k, v)| (*k, *v)).collect(),
         }
@@ -46,7 +44,6 @@ impl From<SszContainer> for ProtoArrayForkChoice {
             ffg_update_required: from.ffg_update_required,
             justified_epoch: from.justified_epoch,
             finalized_epoch: from.finalized_epoch,
-            finalized_root: from.finalized_root,
             nodes: from.nodes,
             indices: HashMap::from_iter(from.indices.into_iter()),
         };

eth2/proto_array_fork_choice/tests/test.rs

@@ -1,4 +1,4 @@
-use lmd_ghost::ProtoArrayForkChoice;
+use proto_array_fork_choice::ProtoArrayForkChoice;
 use types::{Epoch, Hash256};
 
 /// Gives a hash that is not the zero hash (unless i is `usize::max_value)`.
@@ -22,7 +22,7 @@ fn check_bytes_round_trip(original: &ProtoArrayForkChoice) {
 fn no_votes() {
     const VALIDATOR_COUNT: usize = 16;
 
-    let fork_choice = ProtoArrayForkChoice::new(Epoch::new(0), Epoch::new(0), get_hash(0))
+    let fork_choice = ProtoArrayForkChoice::new(Epoch::new(1), Epoch::new(1), get_hash(0))
         .expect("should create fork choice");
 
     check_bytes_round_trip(&fork_choice);
@@ -30,9 +30,9 @@ fn no_votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &[0; VALIDATOR_COUNT]
             )
@@ -47,7 +47,7 @@ fn no_votes() {
     //        /
     //        2
     fork_choice
-        .process_block(get_hash(2), get_hash(0), Epoch::new(0), Epoch::new(0))
+        .process_block(get_hash(2), get_hash(0), Epoch::new(1), Epoch::new(1))
         .expect("should process block");
 
     // Ensure the head is 2
@@ -58,9 +58,9 @@ fn no_votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &[0; VALIDATOR_COUNT]
             )
@@ -75,7 +75,7 @@ fn no_votes() {
     //        / \
     //        2  1
     fork_choice
-        .process_block(get_hash(1), get_hash(0), Epoch::new(0), Epoch::new(0))
+        .process_block(get_hash(1), get_hash(0), Epoch::new(1), Epoch::new(1))
         .expect("should process block");
 
     // Ensure the head is still 2
@@ -86,9 +86,9 @@ fn no_votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &[0; VALIDATOR_COUNT]
             )
@@ -105,7 +105,7 @@ fn no_votes() {
     //           |
     //           3
     fork_choice
-        .process_block(get_hash(3), get_hash(1), Epoch::new(0), Epoch::new(0))
+        .process_block(get_hash(3), get_hash(1), Epoch::new(1), Epoch::new(1))
         .expect("should process block");
 
     // Ensure 3 is the head
@@ -118,9 +118,9 @@ fn no_votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &[0; VALIDATOR_COUNT]
             )
@@ -137,7 +137,7 @@ fn no_votes() {
     //        |  |
     //        4  3
     fork_choice
-        .process_block(get_hash(4), get_hash(2), Epoch::new(0), Epoch::new(0))
+        .process_block(get_hash(4), get_hash(2), Epoch::new(1), Epoch::new(1))
         .expect("should process block");
 
     check_bytes_round_trip(&fork_choice);
@@ -152,9 +152,9 @@ fn no_votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &[0; VALIDATOR_COUNT]
             )
@@ -173,7 +173,7 @@ fn no_votes() {
     //        |
     //        5 <- justified epoch = 1
     fork_choice
-        .process_block(get_hash(5), get_hash(4), Epoch::new(1), Epoch::new(0))
+        .process_block(get_hash(5), get_hash(4), Epoch::new(2), Epoch::new(1))
         .expect("should process block");
 
     // Ensure the head is still 4 whilst the justified epoch is 0.
@@ -188,9 +188,9 @@ fn no_votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &[0; VALIDATOR_COUNT]
             )
@@ -211,9 +211,9 @@ fn no_votes() {
     assert!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 get_hash(5),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &[0; VALIDATOR_COUNT]
             )
@@ -233,9 +233,9 @@ fn no_votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &[0; VALIDATOR_COUNT]
             )
@@ -256,7 +256,7 @@ fn no_votes() {
     //     |
     //     6
     fork_choice
-        .process_block(get_hash(6), get_hash(5), Epoch::new(1), Epoch::new(0))
+        .process_block(get_hash(6), get_hash(5), Epoch::new(2), Epoch::new(1))
         .expect("should process block");
 
     // Ensure 6 is the head
@@ -273,9 +273,9 @@ fn no_votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &[0; VALIDATOR_COUNT]
             )
@@ -291,15 +291,15 @@ fn votes() {
     const VALIDATOR_COUNT: usize = 2;
     let balances = vec![1; VALIDATOR_COUNT];
 
-    let fork_choice = ProtoArrayForkChoice::new(Epoch::new(0), Epoch::new(0), get_hash(0))
+    let fork_choice = ProtoArrayForkChoice::new(Epoch::new(1), Epoch::new(1), get_hash(0))
         .expect("should create fork choice");
 
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &balances
             )
@@ -314,7 +314,7 @@ fn votes() {
     //         /
     //        2
     fork_choice
-        .process_block(get_hash(2), get_hash(0), Epoch::new(0), Epoch::new(0))
+        .process_block(get_hash(2), get_hash(0), Epoch::new(1), Epoch::new(1))
         .expect("should process block");
 
     // Ensure that the head is 2
@@ -325,9 +325,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &balances
             )
@@ -343,7 +343,7 @@ fn votes() {
     //         / \
     //        2   1
     fork_choice
-        .process_block(get_hash(1), get_hash(0), Epoch::new(0), Epoch::new(0))
+        .process_block(get_hash(1), get_hash(0), Epoch::new(1), Epoch::new(1))
         .expect("should process block");
 
     // Ensure that the head is 2
@@ -354,9 +354,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &balances
             )
@@ -371,7 +371,7 @@ fn votes() {
     //         / \
     //        2   1 <- +vote
     fork_choice
-        .process_attestation(0, get_hash(1), Epoch::new(1))
+        .process_attestation(0, get_hash(1), Epoch::new(2))
         .expect("should process attestation");
 
     // Ensure that the head is now 1, beacuse 1 has a vote.
@@ -382,9 +382,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &balances
             )
@@ -399,7 +399,7 @@ fn votes() {
     //          / \
     // +vote-> 2   1
     fork_choice
-        .process_attestation(1, get_hash(2), Epoch::new(1))
+        .process_attestation(1, get_hash(2), Epoch::new(2))
         .expect("should process attestation");
 
     // Ensure that the head is 2 since 1 and 2 both have a vote
@@ -410,9 +410,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &balances
             )
@@ -429,7 +429,7 @@ fn votes() {
     //            |
     //            3
     fork_choice
-        .process_block(get_hash(3), get_hash(1), Epoch::new(0), Epoch::new(0))
+        .process_block(get_hash(3), get_hash(1), Epoch::new(1), Epoch::new(1))
         .expect("should process block");
 
     // Ensure that the head is still 2
@@ -442,9 +442,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &balances
             )
@@ -461,7 +461,7 @@ fn votes() {
     //            |
     //            3 <- +vote
     fork_choice
-        .process_attestation(0, get_hash(3), Epoch::new(2))
+        .process_attestation(0, get_hash(3), Epoch::new(3))
         .expect("should process attestation");
 
     // Ensure that the head is still 2
@@ -474,9 +474,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &balances
             )
@@ -494,7 +494,7 @@ fn votes() {
     //             |
     //             3
     fork_choice
-        .process_attestation(1, get_hash(1), Epoch::new(2))
+        .process_attestation(1, get_hash(1), Epoch::new(3))
         .expect("should process attestation");
 
     // Ensure that the head is now 3
@@ -507,9 +507,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &balances
             )
@@ -528,7 +528,7 @@ fn votes() {
     //            |
     //            4
     fork_choice
-        .process_block(get_hash(4), get_hash(3), Epoch::new(0), Epoch::new(0))
+        .process_block(get_hash(4), get_hash(3), Epoch::new(1), Epoch::new(1))
         .expect("should process block");
 
     // Ensure that the head is now 4
@@ -543,9 +543,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &balances
             )
@@ -566,7 +566,7 @@ fn votes() {
     //           /
     //          5 <- justified epoch = 1
     fork_choice
-        .process_block(get_hash(5), get_hash(4), Epoch::new(1), Epoch::new(1))
+        .process_block(get_hash(5), get_hash(4), Epoch::new(2), Epoch::new(2))
         .expect("should process block");
 
     check_bytes_round_trip(&fork_choice);
@@ -585,9 +585,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &balances
             )
@@ -608,7 +608,7 @@ fn votes() {
     //           / \
     //          5   6 <- justified epoch = 0
     fork_choice
-        .process_block(get_hash(6), get_hash(4), Epoch::new(0), Epoch::new(0))
+        .process_block(get_hash(6), get_hash(4), Epoch::new(1), Epoch::new(1))
         .expect("should process block");
 
     // Move both votes to 5.
@@ -623,10 +623,10 @@ fn votes() {
     //            / \
     // +2 vote-> 5   6
     fork_choice
-        .process_attestation(0, get_hash(5), Epoch::new(3))
+        .process_attestation(0, get_hash(5), Epoch::new(4))
         .expect("should process attestation");
     fork_choice
-        .process_attestation(1, get_hash(5), Epoch::new(3))
+        .process_attestation(1, get_hash(5), Epoch::new(4))
         .expect("should process attestation");
 
     // Add blocks 7, 8 and 9. Adding these blocks helps test the `best_descendant`
@@ -648,13 +648,13 @@ fn votes() {
     //         /
     //         9
     fork_choice
-        .process_block(get_hash(7), get_hash(5), Epoch::new(1), Epoch::new(1))
+        .process_block(get_hash(7), get_hash(5), Epoch::new(2), Epoch::new(2))
         .expect("should process block");
     fork_choice
-        .process_block(get_hash(8), get_hash(7), Epoch::new(1), Epoch::new(1))
+        .process_block(get_hash(8), get_hash(7), Epoch::new(2), Epoch::new(2))
         .expect("should process block");
     fork_choice
-        .process_block(get_hash(9), get_hash(8), Epoch::new(1), Epoch::new(1))
+        .process_block(get_hash(9), get_hash(8), Epoch::new(2), Epoch::new(2))
         .expect("should process block");
 
     // Ensure that 6 is the head, even though 5 has all the votes. This is testing to ensure
@@ -678,9 +678,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
-                Epoch::new(0),
+                Epoch::new(1),
                 Hash256::zero(),
                 &balances
             )
@@ -712,9 +712,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
                 &balances
             )
@@ -744,10 +744,10 @@ fn votes() {
     //         /
     //        9 <- +2 votes
     fork_choice
-        .process_attestation(0, get_hash(9), Epoch::new(4))
+        .process_attestation(0, get_hash(9), Epoch::new(5))
         .expect("should process attestation");
     fork_choice
-        .process_attestation(1, get_hash(9), Epoch::new(4))
+        .process_attestation(1, get_hash(9), Epoch::new(5))
         .expect("should process attestation");
 
     // Add block 10
@@ -768,16 +768,16 @@ fn votes() {
     //         / \
     //        9  10
     fork_choice
-        .process_block(get_hash(10), get_hash(8), Epoch::new(1), Epoch::new(1))
+        .process_block(get_hash(10), get_hash(8), Epoch::new(2), Epoch::new(2))
         .expect("should process block");
 
     // Double-check the head is still 9 (no diagram this time)
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
                 &balances
             )
@@ -807,10 +807,10 @@ fn votes() {
     //         / \
     //        9  10 <- +2 votes
     fork_choice
-        .process_attestation(2, get_hash(10), Epoch::new(4))
+        .process_attestation(2, get_hash(10), Epoch::new(5))
         .expect("should process attestation");
     fork_choice
-        .process_attestation(3, get_hash(10), Epoch::new(4))
+        .process_attestation(3, get_hash(10), Epoch::new(5))
         .expect("should process attestation");
 
     // Check the head is now 10.
@@ -833,9 +833,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
                 &balances
             )
@@ -859,9 +859,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
                 &balances
             )
@@ -885,9 +885,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
                 &balances
             )
@@ -912,9 +912,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
                 &balances
             )
@@ -923,33 +923,39 @@ fn votes() {
         "should find get_hash(9)"
     );
 
-    // Set pruning to an unreachable value.
-    fork_choice.set_prune_threshold(usize::max_value());
-
     check_bytes_round_trip(&fork_choice);
 
-    // Run find-head to trigger a prune.
+    // Set pruning to an unreachable value.
-    assert_eq!(
+    fork_choice.set_prune_threshold(usize::max_value());
-        fork_choice
+    fork_choice
-            .find_head(
+        .update_finalized_root(Epoch::new(2), get_hash(5))
-                Epoch::new(1),
+        .expect("should update finalized root");
-                get_hash(5),
-                Epoch::new(1),
-                get_hash(5),
-                &balances
-            )
-            .expect("should find head"),
-        get_hash(9),
-        "should find get_hash(9)"
-    );
 
     // Ensure that no pruning happened.
     assert_eq!(fork_choice.len(), 11, "there should be 11 blocks");
 
+    // Run find-head
+    assert_eq!(
+        fork_choice
+            .find_head(
+                Epoch::new(2),
+                get_hash(5),
+                Epoch::new(2),
+                get_hash(5),
+                &balances
+            )
+            .expect("should find head"),
+        get_hash(9),
+        "should find get_hash(9)"
+    );
+
     // Set pruning to a value that will result in a prune.
     fork_choice.set_prune_threshold(1);
+    fork_choice
+        .update_finalized_root(Epoch::new(2), get_hash(5))
+        .expect("should update finalized root");
 
-    // Run find-head to trigger a prune.
+    // Run find-head
     //
     //
     //          0
@@ -970,9 +976,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
                 &balances
             )
@@ -998,7 +1004,7 @@ fn votes() {
     //        |
     //        11
     fork_choice
-        .process_block(get_hash(11), get_hash(9), Epoch::new(1), Epoch::new(1))
+        .process_block(get_hash(11), get_hash(9), Epoch::new(2), Epoch::new(2))
         .expect("should process block");
 
     // Ensure the head is now 11
@@ -1015,9 +1021,9 @@ fn votes() {
     assert_eq!(
         fork_choice
             .find_head(
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
-                Epoch::new(1),
+                Epoch::new(2),
                 get_hash(5),
                 &balances
             )