process justifiable for base as well

consensus/fork_choice/src/fork_choice.rs

@@ -603,49 +603,52 @@ where
 
         // Update unrealized justified/finalized checkpoints.
         let (unrealized_justified_checkpoint, unrealized_finalized_checkpoint) = {
-            if !matches!(block, BeaconBlock::Merge(_)) {
-                let (justifiable_beacon_state, _) =
+            let justifiable_beacon_state = match block {
+                BeaconBlock::Merge(_) | BeaconBlock::Altair(_) => {
                     state_processing::per_epoch_processing::altair::process_justifiable(
                         state, spec,
-                    )?;
-                let unrealized_justified_checkpoint =
-                    justifiable_beacon_state.current_justified_checkpoint;
-                let unrealized_finalized_checkpoint = justifiable_beacon_state.finalized_checkpoint;
-
-                // Update best known unrealized justified & finalized checkpoints
-                if unrealized_justified_checkpoint.epoch
-                    > self.fc_store.unrealized_justified_checkpoint().epoch
-                {
-                    self.fc_store
-                        .set_unrealized_justified_checkpoint(unrealized_justified_checkpoint);
+                    )?
+                    .0
                 }
-                if unrealized_finalized_checkpoint.epoch
-                    > self.fc_store.unrealized_finalized_checkpoint().epoch
-                {
-                    self.fc_store
-                        .set_unrealized_finalized_checkpoint(unrealized_finalized_checkpoint);
+                BeaconBlock::Base(_) => {
+                    state_processing::per_epoch_processing::base::process_justifiable(state, spec)?
+                        .0
                 }
+            };
 
-                // If block is from past epochs, try to update store's justified & finalized checkpoints right away
-                if block.slot().epoch(E::slots_per_epoch())
-                    < current_slot.epoch(E::slots_per_epoch())
-                {
-                    self.update_checkpoints(
-                        unrealized_justified_checkpoint,
-                        unrealized_finalized_checkpoint,
-                        state.slot(),
-                        current_slot,
-                        spec,
-                    )?;
-                }
+            let unrealized_justified_checkpoint =
+                justifiable_beacon_state.current_justified_checkpoint;
+            let unrealized_finalized_checkpoint = justifiable_beacon_state.finalized_checkpoint;
 
-                (
-                    Some(unrealized_justified_checkpoint),
-                    Some(unrealized_finalized_checkpoint),
-                )
-            } else {
-                (None, None)
+            // Update best known unrealized justified & finalized checkpoints
+            if unrealized_justified_checkpoint.epoch
+                > self.fc_store.unrealized_justified_checkpoint().epoch
+            {
+                self.fc_store
+                    .set_unrealized_justified_checkpoint(unrealized_justified_checkpoint);
+            }
+            if unrealized_finalized_checkpoint.epoch
+                > self.fc_store.unrealized_finalized_checkpoint().epoch
+            {
+                self.fc_store
+                    .set_unrealized_finalized_checkpoint(unrealized_finalized_checkpoint);
+            }
+
+            // If block is from past epochs, try to update store's justified & finalized checkpoints right away
+            if block.slot().epoch(E::slots_per_epoch()) < current_slot.epoch(E::slots_per_epoch()) {
+                self.update_checkpoints(
+                    unrealized_justified_checkpoint,
+                    unrealized_finalized_checkpoint,
+                    state.slot(),
+                    current_slot,
+                    spec,
+                )?;
             }
+
+            (
+                Some(unrealized_justified_checkpoint),
+                Some(unrealized_finalized_checkpoint),
+            )
         };
 
         let target_slot = block
@@ -695,30 +698,33 @@ where
 
         // This does not apply a vote to the block, it just makes fork choice aware of the block so
         // it can still be identified as the head even if it doesn't have any votes.
-        self.proto_array.process_block(ProtoBlock {
-            slot: block.slot(),
-            root: block_root,
-            parent_root: Some(block.parent_root()),
-            target_root,
-            current_epoch_shuffling_id: AttestationShufflingId::new(
-                block_root,
-                state,
-                RelativeEpoch::Current,
-            )
-            .map_err(Error::BeaconStateError)?,
-            next_epoch_shuffling_id: AttestationShufflingId::new(
-                block_root,
-                state,
-                RelativeEpoch::Next,
-            )
-            .map_err(Error::BeaconStateError)?,
-            state_root: block.state_root(),
-            justified_checkpoint: state.current_justified_checkpoint(),
-            finalized_checkpoint: state.finalized_checkpoint(),
-            execution_status,
-            unrealized_justified_checkpoint,
-            unrealized_finalized_checkpoint,
-        }, current_slot)?;
+        self.proto_array.process_block(
+            ProtoBlock {
+                slot: block.slot(),
+                root: block_root,
+                parent_root: Some(block.parent_root()),
+                target_root,
+                current_epoch_shuffling_id: AttestationShufflingId::new(
+                    block_root,
+                    state,
+                    RelativeEpoch::Current,
+                )
+                .map_err(Error::BeaconStateError)?,
+                next_epoch_shuffling_id: AttestationShufflingId::new(
+                    block_root,
+                    state,
+                    RelativeEpoch::Next,
+                )
+                .map_err(Error::BeaconStateError)?,
+                state_root: block.state_root(),
+                justified_checkpoint: state.current_justified_checkpoint(),
+                finalized_checkpoint: state.finalized_checkpoint(),
+                execution_status,
+                unrealized_justified_checkpoint,
+                unrealized_finalized_checkpoint,
+            },
+            current_slot,
+        )?;
 
         Ok(())
     }

consensus/state_processing/src/per_epoch_processing/base.rs

@@ -7,6 +7,7 @@ use crate::per_epoch_processing::{
 pub use justification_and_finalization::process_justification_and_finalization;
 pub use participation_record_updates::process_participation_record_updates;
 pub use rewards_and_penalties::process_rewards_and_penalties;
+use types::justifiable_beacon_state::JustifiableBeaconState;
 use types::{BeaconState, ChainSpec, EthSpec, RelativeEpoch};
 pub use validator_statuses::{TotalBalances, ValidatorStatus, ValidatorStatuses};
 
@@ -20,18 +21,10 @@ pub fn process_epoch<T: EthSpec>(
     spec: &ChainSpec,
 ) -> Result<EpochProcessingSummary<T>, Error> {
     // Ensure the committee caches are built.
-    state.build_committee_cache(RelativeEpoch::Previous, spec)?;
-    state.build_committee_cache(RelativeEpoch::Current, spec)?;
     state.build_committee_cache(RelativeEpoch::Next, spec)?;
 
-    // Load the struct we use to assign validators into sets based on their participation.
-    //
-    // E.g., attestation in the previous epoch, attested to the head, etc.
-    let mut validator_statuses = ValidatorStatuses::new(state, spec)?;
-    validator_statuses.process_attestations(state)?;
-
-    // Justification and finalization.
-    process_justification_and_finalization(state, &validator_statuses.total_balances, spec)?;
+    let (justifiable_beacon_state, mut validator_statuses) = process_justifiable(state, spec)?;
+    state.update_justifiable(justifiable_beacon_state);
 
     // Rewards and Penalties.
     process_rewards_and_penalties(state, &mut validator_statuses, spec)?;
@@ -72,3 +65,24 @@ pub fn process_epoch<T: EthSpec>(
         statuses: validator_statuses.statuses,
     })
 }
+
+pub fn process_justifiable<T: EthSpec>(
+    state: &mut BeaconState<T>,
+    spec: &ChainSpec,
+) -> Result<(JustifiableBeaconState<T>, ValidatorStatuses), Error> {
+    // Ensure the committee caches are built.
+    state.build_committee_cache(RelativeEpoch::Previous, spec)?;
+    state.build_committee_cache(RelativeEpoch::Current, spec)?;
+
+    // Load the struct we use to assign validators into sets based on their participation.
+    //
+    // E.g., attestation in the previous epoch, attested to the head, etc.
+    let mut validator_statuses = ValidatorStatuses::new(state, spec)?;
+    validator_statuses.process_attestations(state)?;
+
+    // Justification and finalization.
+    let justifiable_beacon_state =
+        process_justification_and_finalization(state, &validator_statuses.total_balances, spec)?;
+
+    Ok((justifiable_beacon_state, validator_statuses))
+}

consensus/state_processing/src/per_epoch_processing/base/justification_and_finalization.rs

@@ -2,25 +2,23 @@ use crate::per_epoch_processing::base::TotalBalances;
 use crate::per_epoch_processing::weigh_justification_and_finalization;
 use crate::per_epoch_processing::Error;
 use safe_arith::SafeArith;
+use types::justifiable_beacon_state::JustifiableBeaconState;
 use types::{BeaconState, ChainSpec, EthSpec};
 
 /// Update the justified and finalized checkpoints for matching target attestations.
 pub fn process_justification_and_finalization<T: EthSpec>(
     state: &mut BeaconState<T>,
     total_balances: &TotalBalances,
     _spec: &ChainSpec,
-) -> Result<(), Error> {
+) -> Result<JustifiableBeaconState<T>, Error> {
     if state.current_epoch() <= T::genesis_epoch().safe_add(1)? {
-        return Ok(());
+        return Ok(JustifiableBeaconState::from(state));
     }
 
-    let justifiable_beacon_state = weigh_justification_and_finalization(
+    weigh_justification_and_finalization(
         state,
         total_balances.current_epoch(),
         total_balances.previous_epoch_target_attesters(),
         total_balances.current_epoch_target_attesters(),
-    )?;
-    state.update_justifiable(justifiable_beacon_state);
-
-    Ok(())
+    )
 }

consensus/types/src/justifiable_beacon_state.rs

@@ -1,9 +1,20 @@
-use crate::{Checkpoint, EthSpec};
+use crate::{BeaconState, Checkpoint, EthSpec};
 use ssz_types::BitVector;
 
-pub struct JustifiableBeaconState<T: EthSpec> {
+pub struct JustifiableBeaconState<E: EthSpec> {
     pub current_justified_checkpoint: Checkpoint,
     pub previous_justified_checkpoint: Checkpoint,
-    pub justification_bits: BitVector<T::JustificationBitsLength>,
+    pub justification_bits: BitVector<E::JustificationBitsLength>,
     pub finalized_checkpoint: Checkpoint,
 }
+
+impl<E: EthSpec> From<&mut BeaconState<E>> for JustifiableBeaconState<E> {
+    fn from(state: &mut BeaconState<E>) -> Self {
+        Self {
+            current_justified_checkpoint: state.current_justified_checkpoint(),
+            previous_justified_checkpoint: state.previous_justified_checkpoint(),
+            justification_bits: state.justification_bits().clone(),
+            finalized_checkpoint: state.finalized_checkpoint(),
+        }
+    }
+}

testing/ef_tests/src/cases/epoch_processing.rs

@@ -89,11 +89,13 @@ impl<E: EthSpec> EpochTransition<E> for JustificationAndFinalization {
             BeaconState::Base(_) => {
                 let mut validator_statuses = base::ValidatorStatuses::new(state, spec)?;
                 validator_statuses.process_attestations(state)?;
-                base::process_justification_and_finalization(
+                let justifiable_beacon_state = base::process_justification_and_finalization(
                     state,
                     &validator_statuses.total_balances,
                     spec,
-                )
+                )?;
+                state.update_justifiable(justifiable_beacon_state);
+                Ok(())
             }
             BeaconState::Altair(_) | BeaconState::Merge(_) => {
                 let prev_participation_cache =