Tidy checkpoint manager

beacon_node/beacon_chain/src/fork_choice.rs

@@ -1,7 +1,7 @@
 mod checkpoint_manager;
 
 use crate::{errors::BeaconChainError, metrics, BeaconChain, BeaconChainTypes};
-use checkpoint_manager::{CheckpointManager, CheckpointWithBalances};
+use checkpoint_manager::{get_effective_balances, CheckpointManager, CheckpointWithBalances};
 use parking_lot::{RwLock, RwLockReadGuard};
 use proto_array_fork_choice::{core::ProtoArray, ProtoArrayForkChoice};
 use ssz_derive::{Decode, Encode};
@@ -59,7 +59,7 @@ impl<T: BeaconChainTypes> ForkChoice<T> {
         let genesis_checkpoint = CheckpointWithBalances {
             epoch: genesis_state.current_epoch(),
             root: genesis_block_root,
-            balances: genesis_state.balances.clone().into(),
+            balances: get_effective_balances(genesis_state),
         };
 
         Self {
@@ -211,7 +211,7 @@ impl<T: BeaconChainTypes> ForkChoice<T> {
         self.backend.maybe_prune(finalized_root).map_err(Into::into)
     }
 
-    /// Returns a read-lock to core `ProtoArray` struct.
+    /// Returns a read-lock to the core `ProtoArray` struct.
     ///
     /// Should only be used when encoding/decoding during troubleshooting.
     pub fn core_proto_array(&self) -> RwLockReadGuard<ProtoArray> {

beacon_node/beacon_chain/src/fork_choice/checkpoint_manager.rs

@@ -33,6 +33,8 @@ impl BalancesCache {
         block_root: Hash256,
         state: &BeaconState<E>,
     ) -> Result<(), Error> {
+        // We are only interested in balances from states that are at the start of an epoch,
+        // because this is where the `current_justified_checkpoint.root` will point.
         if !Self::is_first_block_in_epoch(block_root, state)? {
             return Ok(());
         }
@@ -41,6 +43,8 @@ impl BalancesCache {
         let epoch_boundary_root = if epoch_boundary_slot == state.slot {
             block_root
         } else {
+            // This call remains sensible as long as `state.block_roots` is larger than a single
+            // epoch.
             *state.get_block_root(epoch_boundary_slot)?
         };
 
@@ -62,6 +66,9 @@ impl BalancesCache {
 
     /// Returns `true` if the given `block_root` is the first/only block to have been processed in
     /// the epoch of the given `state`.
+    ///
+    /// We can determine if it is the first block by looking back through `state.block_roots` to
+    /// see if there is a block in the current epoch with a different root.
     fn is_first_block_in_epoch<E: EthSpec>(
         block_root: Hash256,
         state: &BeaconState<E>,
@@ -101,7 +108,7 @@ impl BalancesCache {
 ///
 /// Any validator who is not active in the epoch of the given `state` is assigned a balance of
 /// zero.
-fn get_effective_balances<T: EthSpec>(state: &BeaconState<T>) -> Vec<u64> {
+pub fn get_effective_balances<T: EthSpec>(state: &BeaconState<T>) -> Vec<u64> {
     state
         .validators
         .iter()
@@ -125,8 +132,8 @@ pub struct CheckpointWithBalances {
     /// These are the balances of the state with `self.root`.
     ///
     /// Importantly, these are _not_ the balances of the first state that we saw that has
-    /// `self.epoch` and `self.root` as `state.current_justified_checkpoint`. That state is the
+    /// `self.epoch` and `self.root` as `state.current_justified_checkpoint`. These are the
-    /// state at which the justified state was determined, not the actual justified state.
+    /// balances of the state from the block with `state.current_justified_checkpoint.root`.
     pub balances: Vec<u64>,
 }