perf(beacon): weigh fork-choice votes from a justified-balances snapshot

crates/blockchain/state_transition/src/metrics.rs

@@ -80,6 +80,41 @@ pub fn inc_total_active_balance_lookups(result: &str) {
         .inc();
 }
 
+
+static LEAN_BEACON_JUSTIFIED_BALANCES_LOOKUPS_TOTAL: LazyLock<IntCounterVec> = LazyLock::new(
+    || {
+        register_int_counter_vec!(
+            "lean_beacon_justified_balances_lookups_total",
+            "Beacon fork-choice justified-balances snapshot lookups, by whether the cached snapshot served them",
+            &["result"]
+        )
+        .unwrap()
+    },
+);
+
+/// Count one justified-balances lookup: `hit` (the cached snapshot matched the
+/// justified checkpoint) or `miss` (it was rebuilt from the checkpoint state).
+pub fn inc_justified_balances_lookups(result: &str) {
+    LEAN_BEACON_JUSTIFIED_BALANCES_LOOKUPS_TOTAL
+        .with_label_values(&[result])
+        .inc();
+}
+
+static LEAN_BEACON_JUSTIFIED_BALANCES_BUILD_SECONDS: LazyLock<Histogram> = LazyLock::new(|| {
+    register_histogram!(
+        "lean_beacon_justified_balances_build_seconds",
+        "Duration of one justified-balances snapshot build (one pass over the checkpoint state's registry)",
+        vec![0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0]
+    )
+    .unwrap()
+});
+
+/// Time one justified-balances snapshot build, excluding the checkpoint state
+/// lookup it starts from.
+pub fn time_justified_balances_build() -> TimingGuard {
+    TimingGuard::new(&LEAN_BEACON_JUSTIFIED_BALANCES_BUILD_SECONDS)
+}
+
 static LEAN_STATE_TRANSITION_TIME_SECONDS: LazyLock<Histogram> = LazyLock::new(|| {
     register_histogram!(
         "lean_state_transition_time_seconds",

crates/blockchain/state_transition/tests/beacon_spec/fork_choice.rs

@@ -976,6 +976,108 @@ pub(super) fn apply_checks(
     if let Some(expected) = &checks.get_proposer_head {
         check_get_proposer_head(expected, store, config, committees)?;
     }
+    check_weights_against_the_spec(store, config)?;
+    check_node_weights_against_the_spec(store, config, committees)
+}
+
+/// The gloas counterpart of [`check_weights_against_the_spec`], run at every
+/// `checks` step of a case whose current slot is gloas.
+///
+/// The head walk weighs payload-aware nodes with `compute_node_weights`, whose
+/// vote loop reads the justified-balances snapshot; `gloas_get_weight` is the
+/// specification's per-node definition (payload status, the previous-slot
+/// rule, the gated proposer boost) and reads the justified checkpoint state
+/// itself. Every node under the justified root is compared, not only the
+/// leaves the `viable_for_head_roots_and_weights` check covers.
+fn check_node_weights_against_the_spec(
+    store: &Store,
+    config: &Config,
+    committees: &CommitteeCache,
+) -> Result<(), String> {
+    let current_slot = fork_choice::get_current_slot(store, config);
+    let fork = config.fork_at_epoch(current_slot / preset::SLOTS_PER_EPOCH);
+    if ForkRules::of(fork) != ForkRules::Gloas {
+        return Ok(());
+    }
+    let index = store.block_index();
+    let blocks = fork_choice::get_filtered_block_tree(store, &index, config)
+        .map_err(|err| format!("get_filtered_block_tree: {err:?}"))?;
+    let walked =
+        fork_choice::compute_node_weights(store, &index, config, committees, ForkRules::Gloas)
+            .map_err(|err| format!("compute_node_weights: {err:?}"))?;
+    let mut pending = vec![ForkChoiceNode {
+        root: store.beacon_justified_checkpoint().root,
+        payload_status: PayloadStatus::Pending,
+    }];
+    while let Some(node) = pending.pop() {
+        let spec = match fork_choice::gloas_get_weight(store, node, config, committees) {
+            Ok(spec) => Some(spec),
+            // A vote for a block that invalidation removed from the index: the
+            // spec-literal weight raises on it and the walk drops it on
+            // purpose, so the two are not comparable there.
+            Err(err) if format!("{err:?}").contains("root in store.blocks") => None,
+            Err(err) => {
+                return Err(format!(
+                    "gloas_get_weight(0x{}, {}): {err:?}",
+                    hex::encode(node.root.0),
+                    node.payload_status as u8
+                ));
+            }
+        };
+        if let Some(spec) = spec {
+            let block_slot = index
+                .get(&node.root)
+                .map(|&(slot, _)| slot)
+                .ok_or("a tree node is not in the block index")?;
+            let single_pass = walked.weight(node, block_slot);
+            if spec != single_pass {
+                return Err(format!(
+                    "weight of 0x{} ({}): spec gloas_get_weight {spec}, compute_node_weights {single_pass}",
+                    hex::encode(node.root.0),
+                    node.payload_status as u8
+                ));
+            }
+        }
+        pending.extend(
+            fork_choice::get_node_children(store, &blocks, node)
+                .map_err(|err| format!("get_node_children: {err:?}"))?,
+        );
+    }
+    Ok(())
+}
+
+/// Oracle for the fork-choice weights, run at every `checks` step.
+///
+/// No fixture checks weights directly, but `get_head` descends on
+/// `compute_weights` (one pass over the votes, balances from the justified-
+/// balances snapshot), while `get_weight` is the specification's per-root
+/// definition and reads the justified checkpoint state itself. Comparing the
+/// two for every block in the filtered tree puts the snapshot against the spec
+/// path on every fixture, whatever the fixture asserts.
+fn check_weights_against_the_spec(store: &Store, config: &Config) -> Result<(), String> {
+    let index = store.block_index();
+    let tree = fork_choice::get_filtered_block_tree(store, &index, config)
+        .map_err(|err| format!("get_filtered_block_tree: {err:?}"))?;
+    let weights = fork_choice::compute_weights(store, &index, config)
+        .map_err(|err| format!("compute_weights: {err:?}"))?;
+    for root in tree.keys() {
+        let spec = match fork_choice::get_weight(store, &index, *root, config) {
+            Ok(spec) => spec,
+            // A vote for a block that invalidation removed from the index
+            // (`sync/optimistic`): the specification's `get_weight` raises on
+            // it, and `compute_weights` drops such a vote on purpose, so the
+            // two are not comparable for that case.
+            Err(err) if format!("{err:?}").contains("root in store.blocks") => continue,
+            Err(err) => return Err(format!("get_weight(0x{}): {err:?}", hex::encode(root.0))),
+        };
+        let single_pass = weights.get(root).copied().unwrap_or_default();
+        if spec != single_pass {
+            return Err(format!(
+                "weight of 0x{}: spec get_weight {spec}, compute_weights {single_pass}",
+                hex::encode(root.0)
+            ));
+        }
+    }
     Ok(())
 }
 

crates/common/types/src/beacon/fork_choice.rs

@@ -20,7 +20,10 @@
 
 use libssz_derive::{HashTreeRoot, SszDecode, SszEncode};
 
-use crate::beacon::primitives::{Epoch, ExecutionBlockHash, Root, Slot, Uint256};
+use crate::beacon::containers::Checkpoint;
+use crate::beacon::primitives::{
+    Epoch, ExecutionBlockHash, Gwei, Root, Slot, Uint256, ValidatorIndex,
+};
 
 /// One validator's most recent attestation: the epoch it targeted, and the
 /// block it attested to (the LMD GHOST vote).
@@ -206,12 +209,81 @@ pub struct PayloadStatusV1 {
     pub validation_error: Option<String>,
 }
 
+/// The justified checkpoint state's balances, flattened for the fork-choice
+/// vote loop: `store.checkpoint_states[checkpoint]` as `get_weight` and
+/// `get_proposer_score` read it.
+///
+/// A tree descent per vote (`state.validator(i)`) is an order of magnitude
+/// dearer than an array read, and the loop runs once per `get_head` over every
+/// voter. This is built once per justified checkpoint and keyed by it, so a
+/// reader compares [`Self::checkpoint`] with the store's current one and
+/// rebuilds on a mismatch, with no hook on the code that moves the checkpoint.
+///
+/// Held in the store, hence here rather than in `ethlambda-state-transition`;
+/// the builder lives there, next to the state accessors.
+#[derive(Debug, PartialEq, Eq)]
+pub struct JustifiedBalances {
+    checkpoint: Checkpoint,
+    /// Zero unless the validator is active at `checkpoint.epoch` and unslashed,
+    /// the two conditions under which a vote weighs anything.
+    balances: Box<[Gwei]>,
+    /// The sum of `balances`, floored at one increment: the total the
+    /// committee fraction behind the proposer boost and the reorg thresholds
+    /// divides. Not `get_total_active_balance`, which counts slashed
+    /// validators that are still active.
+    total_balance: Gwei,
+}
+
+impl JustifiedBalances {
+    /// Wraps already-computed values; see the field docs for what they mean.
+    pub fn new(checkpoint: Checkpoint, balances: Box<[Gwei]>, total_balance: Gwei) -> Self {
+        Self {
+            checkpoint,
+            balances,
+            total_balance,
+        }
+    }
+
+    /// The checkpoint these balances were derived from.
+    pub fn checkpoint(&self) -> Checkpoint {
+        self.checkpoint
+    }
+
+    /// The weight of `index`'s vote: zero when inactive or slashed, and also
+    /// past the end, since that validator did not exist at the checkpoint.
+    pub fn get(&self, index: ValidatorIndex) -> Gwei {
+        usize::try_from(index)
+            .ok()
+            .and_then(|index| self.balances.get(index))
+            .copied()
+            .unwrap_or_default()
+    }
+
+    /// The effective balance of the checkpoint state's active, unslashed
+    /// validators, as the specification's `calculate_committee_fraction`
+    /// totals it.
+    pub fn total_balance(&self) -> Gwei {
+        self.total_balance
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use libssz::{SszDecode as _, SszEncode as _};
 
     use super::*;
 
+    #[test]
+    fn justified_balances_read_zero_past_the_registry() {
+        let balances = JustifiedBalances::new(Checkpoint::default(), vec![7, 0, 9].into(), 16);
+        assert_eq!(balances.get(0), 7);
+        assert_eq!(balances.get(1), 0);
+        assert_eq!(balances.get(2), 9);
+        assert_eq!(balances.get(3), 0);
+        assert_eq!(balances.get(u64::MAX), 0);
+        assert_eq!(balances.total_balance(), 16);
+    }
+
     #[test]
     fn a_pow_block_round_trips_through_ssz() {
         // The store persists these, so the derive has to survive the move out