perf(beacon): write epoch balances through an in-order copy-on-write cursor

CLAUDE.md

@@ -589,7 +589,10 @@ existing once per chain.
   is the containers' own round-trip and shape tests.
 - **`Validators` and `Balances` are `ethlambda_ssz_tree::List`s**, persistent
   Merkle trees that cache node hashes and share unchanged subtrees between
-  states through `Arc`; they have no slices and no `iter_mut`. A leaf holds a
+  states through `Arc`; they have no slices and no `iter_mut` (a pass that
+  rewrites most of a list uses the lazy copy-on-write cursor, `iter_cow` /
+  `try_update_each`, with `BeaconState::registry_mut` for disjoint borrows of
+  both lists). A leaf holds a
   page-sized run of elements rather than one chunk, and an inner node a page of
   child pointers spanning several binary levels, so a lookup crosses a handful
   of nodes and a rebuilt leaf or node copies one page. Writes are

crates/blockchain/state_transition/src/beacon/helpers/mutators.rs

@@ -38,6 +38,48 @@ pub fn decrease_balance(state: &mut BeaconState, index: ValidatorIndex, delta: G
     Ok(())
 }
 
+/// Applies reward and penalty vectors to every balance in one in-order pass.
+///
+/// Equivalent to calling [`increase_balance`] then [`decrease_balance`] for
+/// every index, one `(rewards, penalties)` pair after another: each balance
+/// depends on its own deltas only, so going index by index and pair by pair
+/// within an index gives the same result as the specification's pair-by-pair
+/// loop. The saturation stays per step, in spec order (`saturating_add` of the
+/// reward, then `saturating_sub` of the penalty, for each pair): netting the
+/// deltas first would let a reward mask a penalty that should have driven a low
+/// balance to zero.
+///
+/// Goes through the balances' write cursor instead of a `get_mut` per call, so
+/// nothing is buffered per element and a leaf whose balances all come out
+/// unchanged keeps its hash. Every vector is indexed by validator index and
+/// must cover the registry.
+pub fn apply_balance_deltas(
+    state: &mut BeaconState,
+    deltas: &[(Vec<Gwei>, Vec<Gwei>)],
+) -> Result<()> {
+    let validator_count = state.validators().len();
+    if state.balances().len() < validator_count {
+        return Err(Error::UnknownValidator(
+            state.balances().len() as ValidatorIndex
+        ));
+    }
+    state.balances_mut().try_update_each(|balance| {
+        let index = balance.index();
+        // Balances past the registry are not the specification's to touch.
+        if index >= validator_count {
+            return Ok(());
+        }
+        let mut value = **balance;
+        for (rewards, penalties) in deltas {
+            value = value
+                .saturating_add(rewards[index])
+                .saturating_sub(penalties[index]);
+        }
+        balance.set(value);
+        Ok(())
+    })
+}
+
 /// Puts a validator into the exit queue.
 ///
 /// Does nothing if it is already exiting, so this is safe to call more than once

crates/blockchain/state_transition/src/beacon/stf/epoch/altair.rs

@@ -29,9 +29,8 @@ use crate::beacon::helpers::altair::{
 };
 use crate::beacon::helpers::finality::{get_eligible_validator_indices, is_in_inactivity_leak};
 use crate::beacon::helpers::math::saturating_sub;
-use crate::beacon::helpers::mutators::{decrease_balance, increase_balance};
+use crate::beacon::helpers::mutators::apply_balance_deltas;
 use crate::beacon::preset;
-use crate::beacon::primitives::ValidatorIndex;
 
 use super::justification::weigh_justification_and_finalization;
 
@@ -176,11 +175,11 @@ pub fn process_inactivity_updates(state: &mut BeaconState, config: &Config) -> R
 /// [`constants::PARTICIPATION_FLAG_WEIGHTS`]) plus the same kind of inactivity
 /// penalty, computed from [`crate::beacon::helpers::altair::get_inactivity_penalty_deltas`]
 /// against the scores [`process_inactivity_updates`] just brought up to date.
-/// Applying rewards and penalties as two separate passes (through
-/// [`increase_balance`] and [`decrease_balance`], not one netted delta) is
-/// unchanged from phase0, and for the same reason: [`decrease_balance`] floors
-/// at zero, so netting first would let a reward mask a penalty that should
-/// have driven a low balance all the way down.
+/// Rewards and penalties are applied as separate saturating steps, in spec
+/// order (see [`apply_balance_deltas`]), not as one netted delta: that is
+/// unchanged from phase0, and for the same reason: a penalty floors at zero,
+/// so netting first would let a reward mask a penalty that should have driven
+/// a low balance all the way down.
 ///
 /// Skipped entirely at the genesis epoch: rewards pay for participation
 /// recorded during the previous epoch, and genesis has none.
@@ -195,14 +194,7 @@ pub fn process_rewards_and_penalties(state: &mut BeaconState, config: &Config) -
     }
     deltas.push(get_inactivity_penalty_deltas(state, config)?);
 
-    let validator_count = state.validators().len() as ValidatorIndex;
-    for (rewards, penalties) in deltas {
-        for index in 0..validator_count {
-            increase_balance(state, index, rewards[index as usize])?;
-            decrease_balance(state, index, penalties[index as usize])?;
-        }
-    }
-    Ok(())
+    apply_balance_deltas(state, &deltas)
 }
 
 /// Rotates the current epoch's participation flags into the previous slot and
@@ -256,7 +248,7 @@ mod tests {
     use super::*;
     use crate::beacon::fork::ForkName;
     use crate::beacon::helpers::altair::add_flag;
-    use crate::beacon::primitives::BlsPubkey;
+    use crate::beacon::primitives::{BlsPubkey, ValidatorIndex};
 
     /// A deterministic but genuinely valid BLS public key for validator
     /// `index`.

crates/blockchain/state_transition/src/beacon/stf/epoch/electra.rs

@@ -35,7 +35,7 @@ use crate::beacon::bls;
 use crate::beacon::config::Config;
 use crate::beacon::constants::{self, FAR_FUTURE_EPOCH};
 use crate::beacon::containers::shared::{DepositMessage, Validator};
-use crate::beacon::containers::{BeaconState, electra, fulu};
+use crate::beacon::containers::{BeaconState, RegistryMut, electra, fulu};
 use crate::beacon::error::{Error, Result};
 use crate::beacon::helpers::accessors::{get_current_epoch, get_total_active_balance};
 use crate::beacon::helpers::electra::{
@@ -641,28 +641,33 @@ pub fn process_effective_balance_updates(state: &mut BeaconState) -> Result<()>
     const DOWNWARD_THRESHOLD: Gwei = HYSTERESIS_INCREMENT * preset::HYSTERESIS_DOWNWARD_MULTIPLIER;
     const UPWARD_THRESHOLD: Gwei = HYSTERESIS_INCREMENT * preset::HYSTERESIS_UPWARD_MULTIPLIER;
 
-    // Two passes for the same reason `super::process_effective_balance_updates`
-    // needs them: `state` is an enum over per-fork structs, so there is no
-    // way to hold `validators` mutably while also reading `balances`, or
-    // (here) while calling `get_max_effective_balance` on the validator
-    // being decided on.
-    let mut updates = Vec::new();
-    for (index, validator) in state.validators().iter().enumerate() {
-        let balance = state.balances()[index];
+    // Same shape as `super::process_effective_balance_updates`: the registry's
+    // write cursor, with the balances read in step through the disjoint borrow
+    // `registry_mut` hands out. `get_max_effective_balance` reads the validator
+    // being decided on straight from the cursor's element.
+    let RegistryMut {
+        validators,
+        balances,
+    } = state.registry_mut();
+    let mut balances = balances.iter();
+    let mut pass = validators.iter_cow();
+    while let Some(mut validator) = pass.next_cow() {
+        let balance = *balances
+            .next()
+            .expect("balances are positionally parallel to validators");
         if balance + DOWNWARD_THRESHOLD < validator.effective_balance
             || validator.effective_balance + UPWARD_THRESHOLD < balance
         {
-            let max_effective_balance = get_max_effective_balance(validator);
+            let max_effective_balance = get_max_effective_balance(&validator);
             let effective = (balance - balance % preset::EFFECTIVE_BALANCE_INCREMENT)
                 .min(max_effective_balance);
-            updates.push((index, effective));
+            // Equal to the old value when capped at the ceiling: checked before
+            // the copy, so a validator already there dirties nothing.
+            if effective != validator.effective_balance {
+                validator.make_mut().effective_balance = effective;
+            }
         }
     }
-
-    let validators = state.validators_mut();
-    for (index, effective) in updates {
-        validators[index].effective_balance = effective;
-    }
     Ok(())
 }
 

crates/blockchain/state_transition/src/beacon/stf/epoch/mod.rs

@@ -23,7 +23,7 @@ pub mod registry;
 pub mod rewards;
 
 use crate::beacon::containers::phase0::PendingAttestation;
-use crate::beacon::containers::{BeaconState, HistoricalBatch};
+use crate::beacon::containers::{BeaconState, HistoricalBatch, RegistryMut};
 use crate::beacon::error::{Result, verify};
 use crate::beacon::fork::ForkName;
 use crate::beacon::helpers::accessors::{
@@ -265,28 +265,33 @@ pub fn process_effective_balance_updates(state: &mut BeaconState) -> Result<()>
     const DOWNWARD_THRESHOLD: Gwei = HYSTERESIS_INCREMENT * preset::HYSTERESIS_DOWNWARD_MULTIPLIER;
     const UPWARD_THRESHOLD: Gwei = HYSTERESIS_INCREMENT * preset::HYSTERESIS_UPWARD_MULTIPLIER;
 
-    // Decided in one pass and applied in another. The state is an enum over
-    // per-fork structs, so the accessors hand out a borrow of the whole state
-    // rather than of one field, and there is no way to hold `validators` mutably
-    // while reading `balances`. Collecting the decisions first keeps this
-    // fork-independent, which matters because every fork runs this step
-    // unchanged.
-    let mut updates = Vec::new();
-    for (index, validator) in state.validators().iter().enumerate() {
-        let balance = state.balances()[index];
+    // `registry_mut` splits the state into the two lists, so the registry's write
+    // cursor can run while the balances are read in step. It copies a leaf only
+    // for a validator whose effective balance changes and keeps every other leaf
+    // (and its hash) as it was; this stays fork-independent, which matters
+    // because every fork runs this step unchanged.
+    let RegistryMut {
+        validators,
+        balances,
+    } = state.registry_mut();
+    let mut balances = balances.iter();
+    let mut pass = validators.iter_cow();
+    while let Some(mut validator) = pass.next_cow() {
+        let balance = *balances
+            .next()
+            .expect("balances are positionally parallel to validators");
         if balance + DOWNWARD_THRESHOLD < validator.effective_balance
             || validator.effective_balance + UPWARD_THRESHOLD < balance
         {
             let effective = (balance - balance % preset::EFFECTIVE_BALANCE_INCREMENT)
                 .min(preset::MAX_EFFECTIVE_BALANCE);
-            updates.push((index, effective));
+            // Equal to the old value when capped at the maximum: checked before
+            // the copy, so a validator already at its ceiling dirties nothing.
+            if effective != validator.effective_balance {
+                validator.make_mut().effective_balance = effective;
+            }
         }
     }
-
-    let validators = state.validators_mut();
-    for (index, effective) in updates {
-        validators[index].effective_balance = effective;
-    }
     Ok(())
 }
 

crates/blockchain/state_transition/src/beacon/stf/epoch/rewards.rs

@@ -12,9 +12,9 @@
 //! a fixed amount per validator regardless of how many there are.
 //!
 //! [`process_rewards_and_penalties`] applies the results as two separate
-//! passes, rewards then penalties, each through [`increase_balance`] and
-//! [`decrease_balance`] rather than a single netted delta. That distinction is
-//! load-bearing: [`decrease_balance`] floors at zero, and netting the two
+//! saturating steps, rewards then penalties (see [`apply_balance_deltas`]),
+//! rather than a single netted delta. That distinction is
+//! load-bearing: a penalty floors at zero, and netting the two
 //! before applying them would let a reward mask a penalty that should have
 //! driven a low balance all the way down.
 //!
@@ -48,7 +48,7 @@ use crate::beacon::helpers::finality::{
     get_eligible_validator_indices, get_finality_delay, is_in_inactivity_leak,
 };
 use crate::beacon::helpers::math::integer_squareroot;
-use crate::beacon::helpers::mutators::{decrease_balance, increase_balance};
+use crate::beacon::helpers::mutators::apply_balance_deltas;
 use crate::beacon::preset;
 use crate::beacon::primitives::{Gwei, ValidatorIndex};
 
@@ -308,21 +308,17 @@ pub fn get_attestation_deltas(
 ///
 /// Skipped entirely at the genesis epoch: rewards pay for attestations cast in
 /// the previous epoch, and genesis has none. Rewards and penalties are two
-/// separate passes over [`increase_balance`] and [`decrease_balance`], not one
-/// netted delta, because [`decrease_balance`] floors at zero: netting first
+/// separate saturating steps (see [`apply_balance_deltas`]), not one
+/// netted delta, because a penalty floors at zero: netting first
 /// would let a reward mask a penalty that should have driven a low balance all
 /// the way down.
 pub fn process_rewards_and_penalties(state: &mut BeaconState, config: &Config) -> Result<()> {
     if get_current_epoch(state) == constants::GENESIS_EPOCH {
         return Ok(());
     }
 
-    let (rewards, penalties) = get_attestation_deltas(state, config)?;
-    for index in 0..state.validators().len() as ValidatorIndex {
-        increase_balance(state, index, rewards[index as usize])?;
-        decrease_balance(state, index, penalties[index as usize])?;
-    }
-    Ok(())
+    let deltas = [get_attestation_deltas(state, config)?];
+    apply_balance_deltas(state, &deltas)
 }
 
 #[cfg(test)]