Add extra justification change check

eth2/proto_array_fork_choice/src/lib.rs

@@ -7,7 +7,7 @@ use ssz::{Decode, Encode};
 use ssz_container::SszContainer;
 use ssz_derive::{Decode, Encode};
 use std::collections::HashMap;
-use types::{Epoch, Hash256};
+use types::{Epoch, Hash256, Slot};
 
 pub const DEFAULT_PRUNE_THRESHOLD: usize = 256;
 
@@ -96,6 +96,7 @@ impl PartialEq for ProtoArrayForkChoice {
 
 impl ProtoArrayForkChoice {
     pub fn new(
+        finalized_block_slot: Slot,
         justified_epoch: Epoch,
         finalized_epoch: Epoch,
         finalized_root: Hash256,
@@ -110,7 +111,13 @@ impl ProtoArrayForkChoice {
         };
 
         proto_array
-            .on_new_block(finalized_root, None, justified_epoch, finalized_epoch)
+            .on_new_block(
+                finalized_block_slot,
+                finalized_root,
+                None,
+                justified_epoch,
+                finalized_epoch,
+            )
             .map_err(|e| format!("Failed to add finalized block to proto_array: {:?}", e))?;
 
         Ok(Self {
@@ -139,6 +146,7 @@ impl ProtoArrayForkChoice {
 
     pub fn process_block(
         &self,
+        slot: Slot,
         block_root: Hash256,
         parent_root: Hash256,
         justified_epoch: Epoch,
@@ -147,6 +155,7 @@ impl ProtoArrayForkChoice {
         self.proto_array
             .write()
             .on_new_block(
+                slot,
                 block_root,
                 Some(parent_root),
                 justified_epoch,
@@ -160,7 +169,6 @@ impl ProtoArrayForkChoice {
         justified_epoch: Epoch,
         justified_root: Hash256,
         finalized_epoch: Epoch,
-        finalized_root: Hash256,
         justified_state_balances: &[u64],
     ) -> Result<Hash256, String> {
         let mut proto_array = self.proto_array.write();
@@ -211,6 +219,15 @@ impl ProtoArrayForkChoice {
         self.proto_array.read().indices.contains_key(block_root)
     }
 
+    pub fn block_slot(&self, block_root: &Hash256) -> Option<Slot> {
+        let proto_array = self.proto_array.read();
+
+        let i = proto_array.indices.get(block_root)?;
+        let block = proto_array.nodes.get(*i)?;
+
+        Some(block.slot)
+    }
+
     pub fn latest_message(&self, validator_index: usize) -> Option<(Hash256, Epoch)> {
         let votes = self.votes.read();
 

eth2/proto_array_fork_choice/src/proto_array.rs

@@ -1,10 +1,13 @@
 use crate::Error;
 use ssz_derive::{Decode, Encode};
 use std::collections::HashMap;
-use types::{Epoch, Hash256};
+use types::{Epoch, Hash256, Slot};
 
 #[derive(Clone, PartialEq, Debug, Encode, Decode)]
 pub struct ProtoNode {
+    /// The `slot` is not necessary for `ProtoArray`, it just exists so external components can
+    /// easily query the block slot. This is useful for upstream fork choice logic.
+    pub slot: Slot,
     root: Hash256,
     parent: Option<usize>,
     justified_epoch: Epoch,
@@ -131,6 +134,7 @@ impl ProtoArray {
     /// It is only sane to supply a `None` parent for the genesis block.
     pub fn on_new_block(
         &mut self,
+        slot: Slot,
         root: Hash256,
         parent_opt: Option<Hash256>,
         justified_epoch: Epoch,
@@ -139,6 +143,7 @@ impl ProtoArray {
         let node_index = self.nodes.len();
 
         let node = ProtoNode {
+            slot,
             root,
             parent: parent_opt.and_then(|parent| self.indices.get(&parent).copied()),
             justified_epoch,
@@ -376,7 +381,7 @@ impl ProtoArray {
                     // There is no current best-child and the child is viable.
                     change_to_child
                 } else {
-                    // There is no current-best child but the child is not viable.
+                    // There is no current best-child but the child is not viable.
                     no_change
                 }
             };
@@ -388,6 +393,7 @@ impl ProtoArray {
 
         parent.best_child = new_best_child;
         parent.best_descendant = new_best_descendant;
+
         Ok(())
     }