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

CLAUDE.md

@@ -647,7 +647,10 @@ existing once per chain.
 - **`Validators` and `Balances` are `ethlambda_ssz_tree::List`s** (gloas:
   `ethlambda_ssz_tree::ProgressiveList`, same method set), persistent
   Merkle trees that cache node hashes and share unchanged subtrees between
-  states through `Arc`; they have no slices and no `iter_mut`. The registry is
+  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). The registry is
   therefore reached through element-level `BeaconState` accessors
   (`validator`/`validator_mut`, `balance`/`balance_mut`, `push_validator`,
   `iter_validators`/`iter_balances`, `validator_count`, `validators_root`), so

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

@@ -33,6 +33,47 @@ 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.validator_count();
+    let balance_count = state.iter_balances().len();
+    if balance_count < validator_count {
+        return Err(Error::UnknownValidator(balance_count as ValidatorIndex));
+    }
+    state.try_update_balances(|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

@@ -194,11 +194,11 @@ pub(super) fn next_inactivity_score(
 /// [`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 [`crate::beacon::helpers::mutators::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.
@@ -265,10 +265,18 @@ fn apply_rewards_and_penalties(
 
     // The iterator is opaque, so its borrow of `state` lasts until it drops.
     drop(balance_iter);
-    for (index, updated) in changes {
-        *state.balance_mut(index as ValidatorIndex)? = updated;
-    }
-    Ok(())
+    // The changes are in index order, so one cursor pass applies them: nothing
+    // is buffered per change, and the leaves it does not touch keep their hash.
+    let mut changes = changes.into_iter().peekable();
+    state.try_update_balances(|balance| {
+        if let Some(&(index, updated)) = changes.peek()
+            && index == balance.index()
+        {
+            balance.set(updated);
+            changes.next();
+        }
+        Ok::<(), Error>(())
+    })
 }
 
 /// Steps 1-3 of every altair-through-fulu epoch, over one registry scan.

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

@@ -763,37 +763,28 @@ pub fn process_pending_consolidations(state: &mut BeaconState, _config: &Config)
 /// hysteresis test itself is [`super::leaves_hysteresis_band`], shared with
 /// that version.
 pub fn process_effective_balance_updates(state: &mut BeaconState) -> Result<()> {
-    // Two passes for the same reason `super::process_effective_balance_updates`
-    // needs them: `validator_mut` clones the element into the update buffer
-    // on every call, whether or not it is then written (see its own doc), so
-    // deciding and writing in one combined pass would buffer and rehash the
-    // whole registry instead of only the validators that actually move; that
-    // also keeps `get_max_effective_balance`, called here on the validator
-    // being decided on, reading rather than fighting a live mutable borrow.
     // The specification reads `state.balances[index]` for every validator, so
     // a balance list shorter than the registry fails at its first missing
-    // entry. Checked here because the zip below would stop at it silently.
+    // entry. Checked here because the paired pass below would stop at it
+    // silently.
     let balance_count = state.iter_balances().len();
     if balance_count < state.validator_count() {
         return Err(Error::UnknownValidator(balance_count as ValidatorIndex));
     }
-    let mut updates = Vec::new();
-    // Zipped rather than indexed: step 3 only writes the balances that changed,
-    // so `balance(index)` would be a tree descent for most validators.
-    for (index, (validator, balance)) in state
-        .iter_validators()
-        .zip(state.iter_balances())
-        .enumerate()
-    {
+    // Same shape as `super::process_effective_balance_updates`: the registry's
+    // write cursor, with the balances read in step.
+    // `get_max_effective_balance` reads the validator being decided on straight
+    // from the cursor's element.
+    state.try_update_validators_with_balances(|validator, balance| {
         if let Some(effective) = updated_effective_balance(validator, balance)? {
-            updates.push((index as ValidatorIndex, 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;
+            }
         }
-    }
-
-    for (index, effective) in updates {
-        state.validator_mut(index)?.effective_balance = effective;
-    }
-    Ok(())
+        Ok(())
+    })
 }
 
 /// The effective balance `validator` moves to given its `balance`, or `None`

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

@@ -291,39 +291,31 @@ pub fn process_eth1_data_reset(state: &mut BeaconState) -> Result<()> {
 /// and since effective balance feeds the shuffling seed's weighting and every
 /// reward, that would churn far more than it measures.
 pub fn process_effective_balance_updates(state: &mut BeaconState) -> Result<()> {
-    // Decided in one pass and applied in another: `validator_mut` clones the
-    // element into the update buffer on every call, whether or not it is
-    // then written (see its own doc), so deciding and writing in one combined
-    // pass would buffer and rehash the whole registry instead of only the
-    // validators that actually move. Collecting the decisions first also
-    // keeps this fork-independent, which matters because every fork runs
-    // this step unchanged.
     // The specification reads `state.balances[index]` for every validator, so
     // a balance list shorter than the registry fails at its first missing
-    // entry. Checked here because the zip below would stop at it silently.
+    // entry. Checked here because the paired pass below would stop at it
+    // silently.
     let balance_count = state.iter_balances().len();
     if balance_count < state.validator_count() {
         return Err(Error::UnknownValidator(balance_count as ValidatorIndex));
     }
-    let mut updates = Vec::new();
-    // Zipped rather than indexed: step 3 only writes the balances that changed,
-    // so `balance(index)` would be a tree descent for most validators.
-    for (index, (validator, balance)) in state
-        .iter_validators()
-        .zip(state.iter_balances())
-        .enumerate()
-    {
+    // One in-order pass through the registry's write cursor, with the balances
+    // 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.
+    state.try_update_validators_with_balances(|validator, balance| {
         if leaves_hysteresis_band(validator.effective_balance, balance)? {
             let effective = (balance - balance % preset::EFFECTIVE_BALANCE_INCREMENT)
                 .min(preset::MAX_EFFECTIVE_BALANCE);
-            updates.push((index as ValidatorIndex, 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;
+            }
         }
-    }
-
-    for (index, effective) in updates {
-        state.validator_mut(index)?.effective_balance = effective;
-    }
-    Ok(())
+        Ok(())
+    })
 }
 
 /// Whether `balance` has moved far enough from `effective_balance` for the

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.validator_count() 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)]