Strict count unrealized (#3522)

## Issue Addressed

Add a flag that can increase count unrealized strictness, defaults to false

## Proposed Changes

Please list or describe the changes introduced by this PR.

## Additional Info

Please provide any additional information. For example, future considerations
or information useful for reviewers.


Co-authored-by: realbigsean <seananderson33@gmail.com>
Co-authored-by: sean <seananderson33@gmail.com>

consensus/proto_array/src/fork_choice_test_definition.rs

@@ -3,6 +3,7 @@ mod ffg_updates;
 mod no_votes;
 mod votes;
 
+use crate::proto_array::CountUnrealizedFull;
 use crate::proto_array_fork_choice::{Block, ExecutionStatus, ProtoArrayForkChoice};
 use crate::InvalidationOperation;
 use serde_derive::{Deserialize, Serialize};
@@ -87,6 +88,7 @@ impl ForkChoiceTestDefinition {
             junk_shuffling_id.clone(),
             junk_shuffling_id,
             ExecutionStatus::Optimistic(ExecutionBlockHash::zero()),
+            CountUnrealizedFull::default(),
         )
         .expect("should create fork choice struct");
         let equivocating_indices = BTreeSet::new();
@@ -296,8 +298,8 @@ fn get_checkpoint(i: u64) -> Checkpoint {
 
 fn check_bytes_round_trip(original: &ProtoArrayForkChoice) {
     let bytes = original.as_bytes();
-    let decoded =
-        ProtoArrayForkChoice::from_bytes(&bytes).expect("fork choice should decode from bytes");
+    let decoded = ProtoArrayForkChoice::from_bytes(&bytes, CountUnrealizedFull::default())
+        .expect("fork choice should decode from bytes");
     assert!(
         *original == decoded,
         "fork choice should encode and decode without change"

consensus/proto_array/src/lib.rs

@@ -4,7 +4,7 @@ mod proto_array;
 mod proto_array_fork_choice;
 mod ssz_container;
 
-pub use crate::proto_array::InvalidationOperation;
+pub use crate::proto_array::{CountUnrealizedFull, InvalidationOperation};
 pub use crate::proto_array_fork_choice::{Block, ExecutionStatus, ProtoArrayForkChoice};
 pub use error::Error;
 

consensus/proto_array/src/proto_array.rs

@@ -118,6 +118,24 @@ impl Default for ProposerBoost {
     }
 }
 
+/// Indicate whether we should strictly count unrealized justification/finalization votes.
+#[derive(Default, PartialEq, Eq, Debug, Serialize, Deserialize, Copy, Clone)]
+pub enum CountUnrealizedFull {
+    True,
+    #[default]
+    False,
+}
+
+impl From<bool> for CountUnrealizedFull {
+    fn from(b: bool) -> Self {
+        if b {
+            CountUnrealizedFull::True
+        } else {
+            CountUnrealizedFull::False
+        }
+    }
+}
+
 #[derive(PartialEq, Debug, Serialize, Deserialize, Clone)]
 pub struct ProtoArray {
     /// Do not attempt to prune the tree unless it has at least this many nodes. Small prunes
@@ -128,6 +146,7 @@ pub struct ProtoArray {
     pub nodes: Vec<ProtoNode>,
     pub indices: HashMap<Hash256, usize>,
     pub previous_proposer_boost: ProposerBoost,
+    pub count_unrealized_full: CountUnrealizedFull,
 }
 
 impl ProtoArray {
@@ -878,12 +897,14 @@ impl ProtoArray {
             return false;
         }
 
+        let genesis_epoch = Epoch::new(0);
+
         let checkpoint_match_predicate =
             |node_justified_checkpoint: Checkpoint, node_finalized_checkpoint: Checkpoint| {
                 let correct_justified = node_justified_checkpoint == self.justified_checkpoint
-                    || self.justified_checkpoint.epoch == Epoch::new(0);
+                    || self.justified_checkpoint.epoch == genesis_epoch;
                 let correct_finalized = node_finalized_checkpoint == self.finalized_checkpoint
-                    || self.finalized_checkpoint.epoch == Epoch::new(0);
+                    || self.finalized_checkpoint.epoch == genesis_epoch;
                 correct_justified && correct_finalized
             };
 
@@ -898,13 +919,26 @@ impl ProtoArray {
             node.justified_checkpoint,
             node.finalized_checkpoint,
         ) {
-            if node.slot.epoch(E::slots_per_epoch()) < current_slot.epoch(E::slots_per_epoch()) {
-                checkpoint_match_predicate(
-                    unrealized_justified_checkpoint,
-                    unrealized_finalized_checkpoint,
-                )
+            let current_epoch = current_slot.epoch(E::slots_per_epoch());
+
+            // If previous epoch is justified, pull up all tips to at least the previous epoch
+            if CountUnrealizedFull::True == self.count_unrealized_full
+                && (current_epoch > genesis_epoch
+                    && self.justified_checkpoint.epoch + 1 == current_epoch)
+            {
+                unrealized_justified_checkpoint.epoch + 1 >= current_epoch
+            // If previous epoch is not justified, pull up only tips from past epochs up to the current epoch
             } else {
-                checkpoint_match_predicate(justified_checkpoint, finalized_checkpoint)
+                // If block is from a previous epoch, filter using unrealized justification & finalization information
+                if node.slot.epoch(E::slots_per_epoch()) < current_epoch {
+                    checkpoint_match_predicate(
+                        unrealized_justified_checkpoint,
+                        unrealized_finalized_checkpoint,
+                    )
+                // If block is from the current epoch, filter using the head state's justification & finalization information
+                } else {
+                    checkpoint_match_predicate(justified_checkpoint, finalized_checkpoint)
+                }
             }
         } else if let (Some(justified_checkpoint), Some(finalized_checkpoint)) =
             (node.justified_checkpoint, node.finalized_checkpoint)

consensus/proto_array/src/proto_array_fork_choice.rs

@@ -1,4 +1,5 @@
 use crate::error::Error;
+use crate::proto_array::CountUnrealizedFull;
 use crate::proto_array::{
     calculate_proposer_boost, InvalidationOperation, Iter, ProposerBoost, ProtoArray, ProtoNode,
 };
@@ -186,6 +187,7 @@ impl ProtoArrayForkChoice {
         current_epoch_shuffling_id: AttestationShufflingId,
         next_epoch_shuffling_id: AttestationShufflingId,
         execution_status: ExecutionStatus,
+        count_unrealized_full: CountUnrealizedFull,
     ) -> Result<Self, String> {
         let mut proto_array = ProtoArray {
             prune_threshold: DEFAULT_PRUNE_THRESHOLD,
@@ -194,6 +196,7 @@ impl ProtoArrayForkChoice {
             nodes: Vec::with_capacity(1),
             indices: HashMap::with_capacity(1),
             previous_proposer_boost: ProposerBoost::default(),
+            count_unrealized_full,
         };
 
         let block = Block {
@@ -531,8 +534,12 @@ impl ProtoArrayForkChoice {
         SszContainer::from(self).as_ssz_bytes()
     }
 
-    pub fn from_bytes(bytes: &[u8]) -> Result<Self, String> {
+    pub fn from_bytes(
+        bytes: &[u8],
+        count_unrealized_full: CountUnrealizedFull,
+    ) -> Result<Self, String> {
         SszContainer::from_ssz_bytes(bytes)
+            .map(|container| (container, count_unrealized_full))
             .map(Into::into)
             .map_err(|e| format!("Failed to decode ProtoArrayForkChoice: {:?}", e))
     }
@@ -692,6 +699,7 @@ mod test_compute_deltas {
             junk_shuffling_id.clone(),
             junk_shuffling_id.clone(),
             execution_status,
+            CountUnrealizedFull::default(),
         )
         .unwrap();
 

consensus/proto_array/src/ssz_container.rs

@@ -1,6 +1,6 @@
 use crate::proto_array::ProposerBoost;
 use crate::{
-    proto_array::{ProtoArray, ProtoNode},
+    proto_array::{CountUnrealizedFull, ProtoArray, ProtoNode},
     proto_array_fork_choice::{ElasticList, ProtoArrayForkChoice, VoteTracker},
 };
 use ssz::{four_byte_option_impl, Encode};
@@ -41,8 +41,8 @@ impl From<&ProtoArrayForkChoice> for SszContainer {
     }
 }
 
-impl From<SszContainer> for ProtoArrayForkChoice {
-    fn from(from: SszContainer) -> Self {
+impl From<(SszContainer, CountUnrealizedFull)> for ProtoArrayForkChoice {
+    fn from((from, count_unrealized_full): (SszContainer, CountUnrealizedFull)) -> Self {
         let proto_array = ProtoArray {
             prune_threshold: from.prune_threshold,
             justified_checkpoint: from.justified_checkpoint,
@@ -50,6 +50,7 @@ impl From<SszContainer> for ProtoArrayForkChoice {
             nodes: from.nodes,
             indices: from.indices.into_iter().collect::<HashMap<_, _>>(),
             previous_proposer_boost: from.previous_proposer_boost,
+            count_unrealized_full,
         };
 
         Self {