Minor changes to justification code

paulhauner · May 25, 2021 · edfe3d3 · edfe3d3
1 parent 97d1807
commit edfe3d3
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 10 deletions.
diff --git a/consensus/state_processing/src/per_epoch_processing/errors.rs b/consensus/state_processing/src/per_epoch_processing/errors.rs
@@ -21,6 +21,7 @@ pub enum EpochProcessingError {
     SszTypesError(ssz_types::Error),
     ArithError(safe_arith::ArithError),
     InconsistentStateFork(InconsistentFork),
+    InvalidJustificationBit(ssz_types::Error),
 }
 
 impl From<InclusionError> for EpochProcessingError {

diff --git a/consensus/state_processing/src/per_epoch_processing/weigh_justification_and_finalization.rs b/consensus/state_processing/src/per_epoch_processing/weigh_justification_and_finalization.rs
@@ -1,5 +1,6 @@
 use crate::per_epoch_processing::Error;
 use safe_arith::SafeArith;
+use std::ops::Range;
 use types::{BeaconState, Checkpoint, EthSpec};
 
 /// Update the justified and finalized checkpoints for matching target attestations.
@@ -37,29 +38,31 @@ pub fn weigh_justification_and_finalization<T: EthSpec>(
     }
 
     let bits = state.justification_bits().clone();
+    let all_bits_set = |range: Range<usize>| -> Result<bool, Error> {
+        for i in range {
+            if !bits.get(i).map_err(Error::InvalidJustificationBit)? {
+                return Ok(false);
+            }
+        }
+        Ok(true)
+    };
 
     // The 2nd/3rd/4th most recent epochs are all justified, the 2nd using the 4th as source.
-    if (1..4).all(|i| bits.get(i).unwrap_or(false))
-        && old_previous_justified_checkpoint.epoch.safe_add(3)? == current_epoch
+    if all_bits_set(1..4)? && old_previous_justified_checkpoint.epoch.safe_add(3)? == current_epoch
     {
         *state.finalized_checkpoint_mut() = old_previous_justified_checkpoint;
     }
     // The 2nd/3rd most recent epochs are both justified, the 2nd using the 3rd as source.
-    else if (1..3).all(|i| bits.get(i).unwrap_or(false))
-        && old_previous_justified_checkpoint.epoch.safe_add(2)? == current_epoch
+    if all_bits_set(1..3)? && old_previous_justified_checkpoint.epoch.safe_add(2)? == current_epoch
     {
         *state.finalized_checkpoint_mut() = old_previous_justified_checkpoint;
     }
     // The 1st/2nd/3rd most recent epochs are all justified, the 1st using the 3nd as source.
-    if (0..3).all(|i| bits.get(i).unwrap_or(false))
-        && old_current_justified_checkpoint.epoch.safe_add(2)? == current_epoch
-    {
+    if all_bits_set(0..3)? && old_current_justified_checkpoint.epoch.safe_add(2)? == current_epoch {
         *state.finalized_checkpoint_mut() = old_current_justified_checkpoint;
     }
     // The 1st/2nd most recent epochs are both justified, the 1st using the 2nd as source.
-    else if (0..2).all(|i| bits.get(i).unwrap_or(false))
-        && old_current_justified_checkpoint.epoch.safe_add(1)? == current_epoch
-    {
+    if all_bits_set(0..2)? && old_current_justified_checkpoint.epoch.safe_add(1)? == current_epoch {
         *state.finalized_checkpoint_mut() = old_current_justified_checkpoint;
     }