perf(beacon): keep inactivity scores and the ring buffers on the tree

crates/blockchain/state_transition/src/beacon/block_production.rs

@@ -344,6 +344,10 @@ pub fn assemble_block(
         &ExecutionEngine::valid(),
         &CommitteeCache::default(),
     )?;
+    // Fold the block's writes into their trees first, so the root is computed
+    // on (and cached in) the state's own nodes rather than a throwaway copy,
+    // as `state_transition` does before it checks a block's state root.
+    post.apply_pending_mutations();
     block.state_root = post.hash_tree_root();
     Ok(block)
 }

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

@@ -136,8 +136,8 @@ pub fn process_inactivity_updates(state: &mut BeaconState, config: &Config) -> R
     let (_, _, inactivity_scores) = state.altair_validator_lists_mut()?;
     let score_count = inactivity_scores.len();
     for index in eligible_indices {
-        let score = inactivity_scores
-            .get_mut(index as usize)
+        let score = *inactivity_scores
+            .get(index as usize)
             .ok_or(Error::IndexOutOfBounds {
                 index: index as usize,
                 len: score_count,
@@ -146,21 +146,30 @@ pub fn process_inactivity_updates(state: &mut BeaconState, config: &Config) -> R
         // `participating_indices` is ascending and duplicate-free (see
         // `get_unslashed_participating_indices`), so membership is a binary
         // search rather than a linear scan.
-        if participating_indices.binary_search(&index).is_ok() {
+        let mut new_score = if participating_indices.binary_search(&index).is_ok() {
             // `x -= min(1, x)`, written with `saturating_sub` so a
             // already-zero score cannot underflow.
-            *score = saturating_sub(*score, 1);
+            saturating_sub(score, 1)
         } else {
             // The specification treats a `uint64` overflow here as an invalid
             // state rather than a wrapped one, so this is checked rather than
             // left to release-mode wrapping.
-            *score = score.checked_add(config.inactivity_score_bias).ok_or(
-                Error::ArithmeticOverflow("inactivity_scores[index] + INACTIVITY_SCORE_BIAS"),
-            )?;
-        }
+            score
+                .checked_add(config.inactivity_score_bias)
+                .ok_or(Error::ArithmeticOverflow(
+                    "inactivity_scores[index] + INACTIVITY_SCORE_BIAS",
+                ))?
+        };
 
         if !leaking {
-            *score = saturating_sub(*score, config.inactivity_score_recovery_rate);
+            new_score = saturating_sub(new_score, config.inactivity_score_recovery_rate);
+        }
+
+        // Outside a leak almost every score stays zero. A write rebuilds its
+        // leaf even for an equal value, which would unshare the whole list
+        // from the parent state's, so only a changed score is written.
+        if new_score != score {
+            inactivity_scores[index as usize] = new_score;
         }
     }
 
@@ -405,6 +414,57 @@ mod tests {
         assert_eq!(scores[0], 0);
     }
 
+    /// Outside a leak a missed epoch adds the bias and recovery takes it back
+    /// to zero, so every score ends where it started. The pass must then
+    /// leave the list untouched, not just equal: a write of an equal value
+    /// would still unshare the tree from the parent state's.
+    #[test]
+    fn inactivity_updates_that_change_nothing_leave_the_tree_shared() {
+        let mut state = altair_state_with_validators(4);
+        state.apply_pending_mutations();
+        let before = state.altair_validator_lists().unwrap().2.clone();
+
+        process_inactivity_updates(&mut state, &Config::mainnet()).unwrap();
+
+        let (_, _, scores) = state.altair_validator_lists().unwrap();
+        assert!(!scores.has_pending_updates(), "no write was buffered");
+        assert!(scores.ptr_eq(&before));
+        assert_eq!(scores.to_vec(), vec![0; 4]);
+    }
+
+    #[test]
+    fn inactivity_updates_buffer_only_the_scores_that_change() {
+        let mut state = altair_state_with_validators(4);
+        {
+            let (previous_epoch_participation, _, scores) =
+                state.altair_validator_lists_mut().unwrap();
+            scores[2] = 3;
+            // Validator 2 participates, so its score falls to zero; the rest
+            // stay at zero.
+            previous_epoch_participation[2] = add_flag(0, constants::TIMELY_TARGET_FLAG_INDEX);
+        }
+        state.apply_pending_mutations();
+
+        process_inactivity_updates(&mut state, &Config::mainnet()).unwrap();
+
+        let (_, _, scores) = state.altair_validator_lists().unwrap();
+        assert!(scores.has_pending_updates());
+        assert_eq!(scores.to_vec(), vec![0; 4]);
+    }
+
+    #[test]
+    fn inactivity_updates_still_reject_an_overflowing_score() {
+        let mut state = altair_state_with_validators(4);
+        {
+            let (_, _, scores) = state.altair_validator_lists_mut().unwrap();
+            scores[1] = u64::MAX;
+        }
+
+        let result = process_inactivity_updates(&mut state, &Config::mainnet());
+
+        assert!(matches!(result, Err(Error::ArithmeticOverflow(_))));
+    }
+
     // -----------------------------------------------------------------------
     // process_rewards_and_penalties
     // -----------------------------------------------------------------------

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

@@ -317,7 +317,9 @@ pub(crate) fn phase0_state_ref<'a>(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::beacon::fork::ForkName;
     use crate::beacon::helpers::test_state;
+    use crate::beacon::primitives::Root;
 
     #[test]
     fn process_slot_leaves_no_buffered_registry_writes() {
@@ -331,6 +333,10 @@ mod tests {
 
         assert!(!state.balances().has_pending_updates());
         assert!(!state.validators().has_pending_updates());
+        // Its own two roots writes stay buffered until the next flush, which
+        // the following slot (or `process_slots`' last step) performs.
+        assert!(state.state_roots().has_pending_updates());
+        assert!(state.block_roots().has_pending_updates());
     }
 
     /// Epoch processing (`process_rewards_and_penalties` here) writes every
@@ -347,7 +353,109 @@ mod tests {
 
         process_slots(&mut state, target_slot, &config).unwrap();
 
-        assert!(!state.balances().has_pending_updates());
-        assert!(!state.validators().has_pending_updates());
+        assert!(!state.has_pending_mutations());
+    }
+
+    #[test]
+    fn process_slots_flushes_the_roots_it_wrote_in_its_last_slot() {
+        for fork in BEACON_FORKS {
+            let mut state = test_state::with_validators_at(fork, 4);
+            let target_slot = state.slot() + 1;
+
+            process_slots(&mut state, target_slot, &Config::mainnet()).unwrap();
+
+            assert!(!state.has_pending_mutations(), "{fork:?}");
+        }
+    }
+
+    const BEACON_FORKS: [ForkName; 7] = [
+        ForkName::Phase0,
+        ForkName::Altair,
+        ForkName::Bellatrix,
+        ForkName::Capella,
+        ForkName::Deneb,
+        ForkName::Electra,
+        ForkName::Fulu,
+    ];
+
+    /// Writes one element of the field a case names, or answers `false` if the
+    /// fork does not have it.
+    type Write = fn(&mut BeaconState) -> bool;
+
+    fn push_historical_summary(state: &mut BeaconState) -> bool {
+        let summary = containers::HistoricalSummary::default();
+        match state {
+            BeaconState::Capella(s) => s.historical_summaries.push(summary).unwrap(),
+            BeaconState::Deneb(s) => s.historical_summaries.push(summary).unwrap(),
+            BeaconState::Electra(s) => s.historical_summaries.push(summary).unwrap(),
+            BeaconState::Fulu(s) => s.historical_summaries.push(summary).unwrap(),
+            _ => return false,
+        }
+        true
+    }
+
+    /// One write to every tree-backed field of every fork: the flush and the
+    /// pending check must both see all of them, since both come from one
+    /// field list.
+    #[test]
+    fn every_tree_field_of_every_fork_is_flushed_and_checked() {
+        let cases: [(&str, Write); 11] = [
+            ("validators", |s| {
+                s.validator_mut(0).unwrap().effective_balance -= 1;
+                true
+            }),
+            ("balances", |s| {
+                s.balances_mut()[0] += 1;
+                true
+            }),
+            ("block_roots", |s| {
+                s.block_roots_mut()[0] = Root::repeat_byte(1);
+                true
+            }),
+            ("state_roots", |s| {
+                s.state_roots_mut()[0] = Root::repeat_byte(1);
+                true
+            }),
+            ("historical_roots", |s| {
+                s.historical_roots_mut().push(Root::repeat_byte(1)).unwrap();
+                true
+            }),
+            ("eth1_data_votes", |s| {
+                s.eth1_data_votes_mut().push(Default::default()).unwrap();
+                true
+            }),
+            ("randao_mixes", |s| {
+                s.randao_mixes_mut()[0] = Root::repeat_byte(1);
+                true
+            }),
+            ("slashings", |s| {
+                s.slashings_mut()[0] += 1;
+                true
+            }),
+            ("inactivity_scores", |s| {
+                match s.altair_validator_lists_mut() {
+                    Ok((_, _, scores)) => {
+                        scores[0] += 1;
+                        true
+                    }
+                    Err(_) => false,
+                }
+            }),
+            ("historical_summaries", push_historical_summary),
+            // A read-only pass writes nothing and must leave nothing pending.
+            ("none", |_| false),
+        ];
+
+        for fork in BEACON_FORKS {
+            let mut state = test_state::with_validators_at(fork, 4);
+            assert!(!state.has_pending_mutations(), "{fork:?} starts flushed");
+
+            for (name, write) in cases {
+                let wrote = write(&mut state);
+                assert_eq!(state.has_pending_mutations(), wrote, "{fork:?} {name}");
+                state.apply_pending_mutations();
+                assert!(!state.has_pending_mutations(), "{fork:?} {name} flushed");
+            }
+        }
     }
 }

crates/blockchain/state_transition/src/beacon/upgrade.rs

@@ -974,8 +974,8 @@ mod tests {
         // place rather than clearing it, since `historical_summaries` is
         // where new history accumulates from here on.
         assert_eq!(
-            post.historical_roots.into_inner(),
-            bellatrix_state.historical_roots.into_inner()
+            post.historical_roots.to_vec(),
+            bellatrix_state.historical_roots.to_vec()
         );
         assert_eq!(post.next_withdrawal_index, 0);
         assert_eq!(post.next_withdrawal_validator_index, 0);

crates/common/ssz-tree/src/lib.rs

@@ -69,6 +69,19 @@ pub use vector::Vector;
 use libssz::{SszDecode, SszEncode};
 use libssz_merkle::HashTreeRoot;
 
+/// The buffered-write surface of a [`List`] or [`Vector`], independent of its
+/// element type and update map.
+///
+/// Object safe, so a container can hold one `&dyn Buffered` per tree field of
+/// different element types and flush or check them all from one field list.
+pub trait Buffered {
+    /// Folds every buffered write into the tree.
+    fn apply_updates(&mut self);
+
+    /// Whether any write is buffered and not yet folded into the tree.
+    fn has_pending_updates(&self) -> bool;
+}
+
 /// A 32-byte Merkle node.
 pub(crate) type Hash256 = libssz_merkle::Node;
 

crates/common/ssz-tree/src/list.rs

@@ -115,6 +115,16 @@ impl<T: Value, const N: usize, U: UpdateMap<T>> List<T, N, U> {
     }
 }
 
+impl<T: Value, const N: usize, U: UpdateMap<T>> crate::Buffered for List<T, N, U> {
+    fn apply_updates(&mut self) {
+        List::apply_updates(self);
+    }
+
+    fn has_pending_updates(&self) -> bool {
+        List::has_pending_updates(self)
+    }
+}
+
 impl<T: Value, const N: usize, U: UpdateMap<T>> Default for List<T, N, U> {
     fn default() -> Self {
         Self::empty()

crates/common/ssz-tree/src/vector.rs

@@ -87,6 +87,16 @@ impl<T: Value, const N: usize, U: UpdateMap<T>> Vector<T, N, U> {
     }
 }
 
+impl<T: Value, const N: usize, U: UpdateMap<T>> crate::Buffered for Vector<T, N, U> {
+    fn apply_updates(&mut self) {
+        Vector::apply_updates(self);
+    }
+
+    fn has_pending_updates(&self) -> bool {
+        Vector::has_pending_updates(self)
+    }
+}
+
 impl<T: Value + Default, const N: usize, U: UpdateMap<T>> Default for Vector<T, N, U> {
     /// A vector of `N` default-valued elements.
     fn default() -> Self {