diff --git a/CLAUDE.md b/CLAUDE.md index 747989a1..6f8a9b70 100644 --- a/CLAUDE.md +++ b/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 diff --git a/crates/blockchain/state_transition/src/beacon/helpers/mutators.rs b/crates/blockchain/state_transition/src/beacon/helpers/mutators.rs index 6f37c736..75ee3dea 100644 --- a/crates/blockchain/state_transition/src/beacon/helpers/mutators.rs +++ b/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, Vec)], +) -> 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 diff --git a/crates/blockchain/state_transition/src/beacon/stf/epoch/altair.rs b/crates/blockchain/state_transition/src/beacon/stf/epoch/altair.rs index 1b2f0732..6a2867ae 100644 --- a/crates/blockchain/state_transition/src/beacon/stf/epoch/altair.rs +++ b/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`. diff --git a/crates/blockchain/state_transition/src/beacon/stf/epoch/electra.rs b/crates/blockchain/state_transition/src/beacon/stf/epoch/electra.rs index 2f6bf9fc..72695fa4 100644 --- a/crates/blockchain/state_transition/src/beacon/stf/epoch/electra.rs +++ b/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(()) } diff --git a/crates/blockchain/state_transition/src/beacon/stf/epoch/mod.rs b/crates/blockchain/state_transition/src/beacon/stf/epoch/mod.rs index 02695380..7167c4e1 100644 --- a/crates/blockchain/state_transition/src/beacon/stf/epoch/mod.rs +++ b/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(()) } diff --git a/crates/blockchain/state_transition/src/beacon/stf/epoch/rewards.rs b/crates/blockchain/state_transition/src/beacon/stf/epoch/rewards.rs index a6bfab32..09522810 100644 --- a/crates/blockchain/state_transition/src/beacon/stf/epoch/rewards.rs +++ b/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,8 +308,8 @@ 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<()> { @@ -317,12 +317,8 @@ pub fn process_rewards_and_penalties(state: &mut BeaconState, config: &Config) - 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)] diff --git a/crates/common/ssz-tree/src/cursor.rs b/crates/common/ssz-tree/src/cursor.rs new file mode 100644 index 00000000..d27b38fe --- /dev/null +++ b/crates/common/ssz-tree/src/cursor.rs @@ -0,0 +1,526 @@ +//! A lazy copy-on-write cursor: one in-order pass that can rewrite any element. +//! +//! Writing most of a list through `get_mut` pays for a pending-write map (an +//! entry per element, then a sorted copy of it) and for a descent per element, +//! only to rebuild each touched leaf in the end. The cursor skips the map. It +//! walks the leaves left to right and copies a leaf only when an element in it +//! is first written; when it leaves the leaf it compares what was written with +//! the original, and keeps the original (its `Arc` and its cached hash) if +//! nothing differs. The leaves that did change are swapped into the tree in +//! one descent when the cursor is dropped, which rebuilds only the paths to +//! them. +//! +//! The tree is never modified while the pass runs, so nothing is buffered +//! between elements: an early return or a panic still leaves a consistent +//! list holding the writes made so far. + +use std::ops::Deref; +use std::sync::Arc; + +use crate::Value; +use crate::tree::{LeafReplacements, Tree, element_root}; + +/// The leaf the cursor is in: the original, plus a copy made on the first +/// write. +struct LeafState { + /// The original leaf, always a [`Tree::Leaf`] while set. + leaf: Option>>, + /// Index of the leaf's first element. + first: usize, + /// Elements in the leaf. + len: usize, + /// Offset of the next element to hand out. + next: usize, + /// The leaf's elements with this pass's writes, made on first write. + copy: Option>, + /// One bit per offset that `make_mut` handed out, so the end-of-leaf + /// comparison looks only at those. All zero whenever `copy` is `None`. + touched: Vec, +} + +impl LeafState { + fn empty() -> Self { + Self { + leaf: None, + first: 0, + len: 0, + next: 0, + copy: None, + touched: Vec::new(), + } + } + + fn original(&self) -> &[T] { + match self.leaf.as_deref() { + Some(Tree::Leaf(leaf)) => &leaf.values, + _ => &[], + } + } + + /// The leaf's elements as the pass sees them: the copy once there is one. + fn current(&self) -> &[T] { + self.copy.as_deref().unwrap_or_else(|| self.original()) + } + + fn load(&mut self, leaf: Arc>, first: usize) { + let Tree::Leaf(inner) = &*leaf else { + unreachable!("the walk yields leaves") + }; + self.len = inner.values.len(); + self.leaf = Some(leaf); + self.first = first; + self.next = 0; + } + + fn make_mut(&mut self, offset: usize) -> &mut T { + if self.copy.is_none() { + self.copy = Some(self.original().to_vec()); + self.touched.resize(self.len.div_ceil(64), 0); + } + self.touched[offset / 64] |= 1 << (offset % 64); + &mut self.copy.as_mut().expect("made just above")[offset] + } + + /// Leaves the current leaf behind: the replacement to swap into the tree, + /// or `None` if the pass left it as it was. + /// + /// A kept original and a carried element root both follow a comparison + /// here, never an assumption that a write changed (or did not change) + /// something. + fn finish(&mut self) -> Option<(usize, Arc>)> { + let copy = self.copy.take()?; + let touched = std::mem::take(&mut self.touched); + let Some(Tree::Leaf(original)) = self.leaf.as_deref() else { + unreachable!("a copy is only made of a loaded leaf") + }; + // The element roots the original already has: the copy carries them + // over, with the changed elements' recomputed. Only a composite leaf + // that was hashed has any. + let mut roots = original.roots.get().map(|roots| roots.to_vec()); + let mut changed = false; + for (word_index, &word) in touched.iter().enumerate() { + let mut bits = word; + while bits != 0 { + let offset = word_index * 64 + bits.trailing_zeros() as usize; + bits &= bits - 1; + if copy[offset] == original.values[offset] { + continue; + } + changed = true; + if let Some(roots) = roots.as_mut() { + roots[offset] = element_root(©[offset]); + } + } + } + // Hand the zeroed bitset back for the next copy. + self.touched = touched; + self.touched.fill(0); + if !changed { + return None; + } + let leaf = Tree::leaf(copy, roots.map(Vec::into_boxed_slice)); + Some((self.first, Arc::new(leaf))) + } +} + +/// An in-order pass over a [`List`](crate::List) or [`Vector`](crate::Vector) +/// that can rewrite any element, made by `iter_cow`. +/// +/// It is lending: [`IterCow::next_cow`] returns an [`ElemCow`] that borrows +/// the pass until the next call, so it is used in a `while let` loop rather +/// than a `for`. The writes reach the list when the pass is dropped, which +/// also happens on an early return or a panic. +pub struct IterCow<'a, T: Value> { + tree: &'a mut Arc>, + depth: usize, + /// For each inner node on the path to the next leaf, the index of its next + /// child to visit. + stack: Vec<(Arc>, usize)>, + /// A tree no taller than a leaf is one, yielded first. + root_leaf: Option>>, + /// Index of the first element of the next leaf to visit. + next_first: usize, + state: LeafState, + /// Leaves that changed, in order, with their first index. + replaced: Vec<(usize, Arc>)>, +} + +impl<'a, T: Value> IterCow<'a, T> { + /// A pass over the tree of height `depth`. The caller has applied every + /// pending write, so the tree is the whole list. + pub(crate) fn new(tree: &'a mut Arc>, depth: usize) -> Self { + let root = Arc::clone(tree); + let mut stack = Vec::new(); + let mut root_leaf = None; + match &*root { + Tree::Node(_) => stack.push((root, 0)), + Tree::Leaf(_) => root_leaf = Some(root), + Tree::Zero(_) => {} + } + Self { + tree, + depth, + stack, + root_leaf, + next_first: 0, + state: LeafState::empty(), + replaced: Vec::new(), + } + } + + /// The next element, to read or to write, or `None` after the last. + pub fn next_cow(&mut self) -> Option> { + loop { + if self.state.next < self.state.len { + let offset = self.state.next; + self.state.next += 1; + let index = self.state.first + offset; + return Some(ElemCow { + state: &mut self.state, + offset, + index, + }); + } + self.finish_leaf(); + let leaf = self.next_leaf()?; + let first = self.next_first; + self.next_first += crate::packing_factor::() << crate::leaf_height::(self.depth); + self.state.load(leaf, first); + } + } + + fn finish_leaf(&mut self) { + if let Some(replacement) = self.state.finish() { + self.replaced.push(replacement); + } + self.state.len = 0; + self.state.next = 0; + } + + /// The next leaf in order, or `None` when the data ends. + fn next_leaf(&mut self) -> Option>> { + if let Some(leaf) = self.root_leaf.take() { + return Some(leaf); + } + loop { + let (node, slot) = self.stack.last_mut()?; + let Tree::Node(inner) = &**node else { + unreachable!("only inner nodes are on the stack") + }; + let Some(child) = inner.children.get(*slot).map(Arc::clone) else { + self.stack.pop(); + continue; + }; + *slot += 1; + match &*child { + Tree::Node(_) => self.stack.push((child, 0)), + Tree::Leaf(_) => return Some(child), + // The data is a prefix: everything from here on is padding. + Tree::Zero(_) => { + self.stack.clear(); + return None; + } + } + } + } +} + +impl Drop for IterCow<'_, T> { + /// Finishes the current leaf and swaps the changed leaves into the tree, + /// rebuilding only the paths to them. + fn drop(&mut self) { + self.finish_leaf(); + if self.replaced.is_empty() { + return; + } + let replaced = std::mem::take(&mut self.replaced); + let mut source = LeafReplacements::new(replaced); + *self.tree = Tree::with_rebuilt_leaves(&*self.tree, self.depth, 0, &mut source); + } +} + +/// One element of an [`IterCow`] pass. +/// +/// Dereferences to the element, with this pass's writes to it. The leaf holding +/// it is copied only by the first [`make_mut`](ElemCow::make_mut) or changing +/// [`set`](ElemCow::set) on any element in it. +pub struct ElemCow<'c, T: Value> { + state: &'c mut LeafState, + offset: usize, + index: usize, +} + +impl Deref for ElemCow<'_, T> { + type Target = T; + + fn deref(&self) -> &T { + &self.state.current()[self.offset] + } +} + +impl ElemCow<'_, T> { + /// The element's index in the list. + pub fn index(&self) -> usize { + self.index + } + + /// The element, to be written. Copies its leaf on the first write to the + /// leaf in this pass; the leaf is kept as it was if the writes in the end + /// leave every element equal. + /// + /// Prefer [`set`](ElemCow::set), or check first, when most writes change + /// nothing: the copy is paid on the first call whether or not the value + /// ends up different. + pub fn make_mut(&mut self) -> &mut T { + self.state.make_mut(self.offset) + } + + /// Writes `value` if it differs from the element, so a write of the same + /// value copies nothing. + pub fn set(&mut self, value: T) { + if **self != value { + *self.make_mut() = value; + } + } +} + +#[cfg(test)] +mod tests { + use std::panic::{AssertUnwindSafe, catch_unwind}; + + use libssz_merkle::{HashTreeRoot, Sha2Hasher, merkleize, pack}; + + use super::*; + use crate::{Hash256, List}; + + fn children(tree: &Tree) -> &[Arc>] { + match tree { + Tree::Node(inner) => &inner.children, + _ => panic!("not an inner node"), + } + } + + fn u64_root(values: &[u64], depth: usize) -> Hash256 { + let bytes: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + merkleize(&Sha2Hasher, &pack(&bytes), Some(1 << depth)) + } + + fn list_root(value: &V) -> [u8; 32] { + HashTreeRoot::hash_tree_root(value, &Sha2Hasher) + } + + /// A tree of four u64 leaves (512 to a leaf) under one inner node. + const DEPTH: usize = 9; + + #[test] + fn only_the_changed_leafs_path_is_new() { + let values: Vec = (0..2048).collect(); + let (mut tree, _) = Tree::from_values(values.clone(), DEPTH); + let before = Arc::clone(&tree); + before.hash(DEPTH); + { + let mut pass = IterCow::new(&mut tree, DEPTH); + while let Some(mut element) = pass.next_cow() { + match element.index() { + // Changes leaf 1. + 600 => element.set(0), + // Touches leaf 2 without changing it. + 1100 => *element.make_mut() = 1100, + // Sets the same value in leaf 3. + 1600 => element.set(1600), + _ => {} + } + } + } + let (old, new) = (children(&before), children(&tree)); + assert!(Arc::ptr_eq(&old[0], &new[0])); + assert!(!Arc::ptr_eq(&old[1], &new[1])); + assert!(Arc::ptr_eq(&old[2], &new[2])); + assert!(Arc::ptr_eq(&old[3], &new[3])); + // The kept leaves keep their cached hash; the new one has none yet. + assert!(new[0].cached_hash().is_some()); + assert!(new[1].cached_hash().is_none()); + + let mut expected = values; + expected[600] = 0; + assert_eq!(tree.hash(DEPTH), u64_root(&expected, DEPTH)); + // The original is untouched. + assert_eq!(before.get(600, DEPTH), Some(&600)); + } + + #[test] + fn a_pass_that_changes_nothing_keeps_the_root_node() { + let (mut tree, _) = Tree::from_values((0..2048u64).collect::>(), DEPTH); + let before = Arc::clone(&tree); + { + let mut pass = IterCow::new(&mut tree, DEPTH); + while let Some(mut element) = pass.next_cow() { + let same = *element; + element.set(same); + *element.make_mut() = same; + } + } + assert!(Arc::ptr_eq(&before, &tree)); + } + + fn composite(n: u8) -> [u8; 48] { + [n; 48] + } + + #[test] + fn carried_element_roots_equal_fresh_ones() { + let values: Vec<[u8; 48]> = (0..200u8).map(composite).collect(); + let (mut tree, _) = Tree::from_values(values.clone(), 8); + tree.hash(8); + { + let mut pass = IterCow::new(&mut tree, 8); + while let Some(mut element) = pass.next_cow() { + match element.index() { + 70 => element.set(composite(250)), + // Touched but equal: its root must stay the old one. + 71 => *element.make_mut() = composite(71), + _ => {} + } + } + } + let mut expected = values; + expected[70] = composite(250); + let Tree::Leaf(leaf) = &*children(&tree)[1] else { + panic!("a 48-byte element's leaf height is 6"); + }; + let carried = leaf.roots.get().expect("carried from the hashed leaf"); + let fresh: Vec = expected[64..128].iter().map(element_root).collect(); + assert_eq!(&carried[..], &fresh[..]); + let roots: Vec = expected.iter().map(element_root).collect(); + assert_eq!(tree.hash(8), merkleize(&Sha2Hasher, &roots, Some(256))); + } + + #[test] + fn a_leaf_without_roots_gets_none_and_hashes_correctly() { + let values: Vec<[u8; 48]> = (0..70u8).map(composite).collect(); + let (mut tree, _) = Tree::from_values(values.clone(), 8); + { + let mut pass = IterCow::new(&mut tree, 8); + while let Some(mut element) = pass.next_cow() { + if element.index() == 3 { + element.set(composite(200)); + } + } + } + let Tree::Leaf(leaf) = &*children(&tree)[0] else { + panic!("expected a leaf"); + }; + assert!(leaf.roots.get().is_none()); + let mut expected = values; + expected[3] = composite(200); + let roots: Vec = expected.iter().map(element_root).collect(); + assert_eq!(tree.hash(8), merkleize(&Sha2Hasher, &roots, Some(256))); + } + + #[test] + fn an_early_return_keeps_the_writes_made_so_far() { + let mut list = List::::try_from((0..1500u64).collect::>()).unwrap(); + let result = list.try_update_each(|element| { + if element.index() == 700 { + return Err("stop"); + } + element.set(element.index() as u64 + 1); + Ok(()) + }); + assert_eq!(result, Err("stop")); + let expected: Vec = (0..1500u64) + .map(|i| if i < 700 { i + 1 } else { i }) + .collect(); + assert_eq!(list.to_vec(), expected); + let reference = List::::try_from(expected).unwrap(); + assert_eq!(list_root(&list), list_root(&reference)); + } + + #[test] + fn a_panic_mid_pass_leaves_a_consistent_list() { + let mut list = List::::try_from((0..1500u64).collect::>()).unwrap(); + let outcome = catch_unwind(AssertUnwindSafe(|| { + let mut pass = list.iter_cow(); + while let Some(mut element) = pass.next_cow() { + if element.index() == 800 { + *element.make_mut() = 9999; + panic!("mid-pass"); + } + element.set(element.index() as u64 + 1); + } + })); + assert!(outcome.is_err()); + let expected: Vec = (0..1500u64) + .map(|i| match i { + 800 => 9999, + i if i < 800 => i + 1, + i => i, + }) + .collect(); + assert_eq!(list.to_vec(), expected); + let reference = List::::try_from(expected).unwrap(); + assert_eq!(list_root(&list), list_root(&reference)); + } + + #[test] + fn empty_lists_and_partial_last_leaves() { + let mut empty = List::::empty(); + assert!(empty.iter_cow().next_cow().is_none()); + empty + .try_update_each(|_| -> Result<(), ()> { unreachable!("no elements") }) + .unwrap(); + + // 513 u64s: a full leaf and a one-element leaf. + let mut list = List::::try_from((0..513u64).collect::>()).unwrap(); + list.try_update_each(|e| -> Result<(), ()> { + e.set(**e + 1); + Ok(()) + }) + .unwrap(); + assert_eq!(list.to_vec(), (1..514u64).collect::>()); + } + + #[test] + fn a_registry_depth_list_updates_across_inner_nodes() { + let boundary = 512 * 512; + let len = boundary + 1000; + let values: Vec = (0..len as u64).collect(); + let mut list = List::::try_from(values.clone()).unwrap(); + list.try_update_each(|e| -> Result<(), ()> { + let index = e.index(); + if index == 5 || index == boundary - 1 || index == boundary || index == len - 1 { + e.set(0); + } + Ok(()) + }) + .unwrap(); + let mut expected = values; + for index in [5, boundary - 1, boundary, len - 1] { + expected[index] = 0; + } + let reference = List::::try_from(expected.clone()).unwrap(); + assert_eq!(list_root(&list), list_root(&reference)); + assert_eq!(list.to_vec(), expected); + } + + #[test] + fn pending_writes_are_applied_before_the_pass() { + let mut list = List::::try_from((0..1000u64).collect::>()).unwrap(); + list.push(1000).unwrap(); + *list.get_mut(3).unwrap() = 77; + list.try_update_each(|e| -> Result<(), ()> { + if e.index() == 3 { + assert_eq!(**e, 77); + } + if e.index() == 1000 { + e.set(5); + } + Ok(()) + }) + .unwrap(); + assert!(!list.has_pending_updates()); + assert_eq!(list[3], 77); + assert_eq!(list[1000], 5); + assert_eq!(list.len(), 1001); + } +} diff --git a/crates/common/ssz-tree/src/interface.rs b/crates/common/ssz-tree/src/interface.rs index 671f6c60..19f3b8ff 100644 --- a/crates/common/ssz-tree/src/interface.rs +++ b/crates/common/ssz-tree/src/interface.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use libssz::{BYTES_PER_LENGTH_OFFSET, DecodeError, SszDecode, SszEncode}; +use crate::cursor::IterCow; use crate::iter::{Iter, TreeIter}; use crate::tree::Tree; use crate::update_map::UpdateMap; @@ -93,6 +94,16 @@ impl> Interface { self.committed_len = self.len; } + /// Applies pending writes, then starts an in-order pass whose writes reach + /// the tree when it is dropped, bypassing the pending-write map. + /// + /// Applying first keeps the pass simple: it reads the tree alone, so a + /// write made earlier is never hidden by (or lost to) the pass. + pub(crate) fn iter_cow(&mut self) -> IterCow<'_, T> { + self.apply_updates(); + IterCow::new(&mut self.tree, self.depth) + } + pub(crate) fn iter_from(&self, start: usize) -> Iter<'_, T, U> { let start = start.min(self.len); Iter { diff --git a/crates/common/ssz-tree/src/lib.rs b/crates/common/ssz-tree/src/lib.rs index 561be2a9..9987cee2 100644 --- a/crates/common/ssz-tree/src/lib.rs +++ b/crates/common/ssz-tree/src/lib.rs @@ -46,6 +46,16 @@ //! a throwaway copy, so none of the new hashes are kept. Call `apply_updates` //! before hashing on a hot path. //! +//! # Rewriting most of a list: the cursor +//! +//! A pass that writes most elements, such as an epoch's balance update, should +//! not go through the buffer: it would hold an entry per element and then +//! sort them. `iter_cow` walks the list in order instead and copies a leaf +//! only when a write in it changes something. A leaf the pass leaves as it was +//! keeps its `Arc` and its hash, and a rebuilt composite leaf keeps the roots of +//! the elements that did not change. The changed leaves are swapped in once, +//! when the pass is dropped. Pending writes are applied before the pass starts. +//! //! # The hasher argument //! //! `HashTreeRoot::hash_tree_root` takes a hasher; these types ignore it and @@ -53,6 +63,7 @@ //! valid for the function that produced it, and SHA-256 is the only one the //! consensus specs use. +mod cursor; mod interface; mod iter; mod list; @@ -61,6 +72,7 @@ mod tree; mod update_map; mod vector; +pub use cursor::{ElemCow, IterCow}; pub use iter::Iter; pub use list::List; pub use update_map::{UpdateMap, VecMap}; diff --git a/crates/common/ssz-tree/src/list.rs b/crates/common/ssz-tree/src/list.rs index f0f98b6f..62efc6ab 100644 --- a/crates/common/ssz-tree/src/list.rs +++ b/crates/common/ssz-tree/src/list.rs @@ -7,6 +7,7 @@ use libssz::{DecodeError, SszDecode, SszEncode}; use libssz_merkle::{HashTreeRoot, Sha2Hasher, Sha256Hasher, mix_in_length}; use libssz_types::TypeError; +use crate::cursor::{ElemCow, IterCow}; use crate::interface::Interface; use crate::iter::Iter; use crate::update_map::{UpdateMap, VecMap}; @@ -76,6 +77,27 @@ impl> List { Ok(()) } + /// An in-order pass that can rewrite any element, for updating most of the + /// list at once. Applies pending writes first; the pass's writes reach the + /// tree when it is dropped, so an early return or a panic keeps the writes + /// made so far. + pub fn iter_cow(&mut self) -> IterCow<'_, T> { + self.interface.iter_cow() + } + + /// Runs `f` on every element in order, stopping at the first error and + /// keeping the writes made before it. See [`List::iter_cow`]. + pub fn try_update_each( + &mut self, + mut f: impl FnMut(&mut ElemCow<'_, T>) -> Result<(), E>, + ) -> Result<(), E> { + let mut pass = self.iter_cow(); + while let Some(mut element) = pass.next_cow() { + f(&mut element)?; + } + Ok(()) + } + /// The elements in order, pending writes included. pub fn iter(&self) -> Iter<'_, T, U> { self.interface.iter_from(0) diff --git a/crates/common/ssz-tree/src/tree.rs b/crates/common/ssz-tree/src/tree.rs index 2eb8c70d..9339e722 100644 --- a/crates/common/ssz-tree/src/tree.rs +++ b/crates/common/ssz-tree/src/tree.rs @@ -86,7 +86,7 @@ impl Tree { }) } - fn leaf(values: Vec, roots: Option>) -> Self { + pub(crate) fn leaf(values: Vec, roots: Option>) -> Self { Tree::Leaf(Leaf { hash: OnceLock::new(), roots: roots.map(OnceLock::from).unwrap_or_default(), @@ -272,16 +272,32 @@ impl Tree { where I: Iterator, { + Self::with_rebuilt_leaves(node, height, first, &mut ElementUpdates(updates)) + } + + /// A copy of `node` whose leaves are rebuilt by `source`, sharing every + /// child the source does not touch. + /// + /// The descent behind [`Tree::with_updated_leaves`], generalized over what + /// happens at a leaf: `source` names, in ascending order, the next index it + /// has work for, and builds the new leaf. Everything above the leaves is + /// rebuilt only along those paths. + pub(crate) fn with_rebuilt_leaves>( + node: &Arc, + height: usize, + first: usize, + source: &mut S, + ) -> Arc { let packing = packing_factor::(); let end = first + (packing << height); - match updates.peek() { - Some(&(index, _)) if index < end => {} + match source.peek() { + Some(index) if index < end => {} _ => return Arc::clone(node), } // A tree shorter than a leaf is one leaf, at its root; otherwise the // walk meets the leaves at the leaf height itself. if height <= max_leaf_height::() { - return Arc::new(node.updated_leaf(first, end, updates)); + return source.leaf(node, first, end); } let below = child_height::(height); let span = packing << below; @@ -290,24 +306,24 @@ impl Tree { Tree::Zero(_) => Vec::new(), Tree::Leaf(_) => unreachable!("a leaf above the leaf height"), }; - while let Some(&(index, _)) = updates.peek() { + while let Some(index) = source.peek() { if index >= end { break; } let slot = (index - first) / span; let child_first = first + slot * span; if let Some(child) = children.get(slot) { - children[slot] = Self::with_updated_leaves(child, below, child_first, updates); + children[slot] = Self::with_rebuilt_leaves(child, below, child_first, source); } else { // The data is a prefix, so a child past the last one only // appears as the next one, grown from nothing by pushes. assert_eq!(slot, children.len(), "a new child follows the last one"); let empty = Arc::new(Tree::Zero(below)); - children.push(Self::with_updated_leaves( + children.push(Self::with_rebuilt_leaves( &empty, below, child_first, - updates, + source, )); } } @@ -356,6 +372,53 @@ impl Tree { } } +/// What a descent does at the leaves it reaches: see +/// [`Tree::with_rebuilt_leaves`]. +pub(crate) trait LeafSource { + /// The index of the next element or leaf this source has work for, in + /// strictly ascending order, or `None` when it is done. + fn peek(&mut self) -> Option; + + /// The replacement for `node`, the leaf covering elements `first..end`, + /// consuming the work inside that range. + fn leaf(&mut self, node: &Arc>, first: usize, end: usize) -> Arc>; +} + +/// Element-level writes, folded into each leaf they land in. +struct ElementUpdates<'a, I: Iterator>(&'a mut Peekable); + +impl> LeafSource for ElementUpdates<'_, I> { + fn peek(&mut self) -> Option { + self.0.peek().map(|&(index, _)| index) + } + + fn leaf(&mut self, node: &Arc>, first: usize, end: usize) -> Arc> { + Arc::new(node.updated_leaf(first, end, self.0)) + } +} + +/// Whole leaves built elsewhere, to swap in: `(first index, new leaf)` in +/// ascending order. +pub(crate) struct LeafReplacements(Peekable>)>>); + +impl LeafReplacements { + pub(crate) fn new(replacements: Vec<(usize, Arc>)>) -> Self { + Self(replacements.into_iter().peekable()) + } +} + +impl LeafSource for LeafReplacements { + fn peek(&mut self) -> Option { + self.0.peek().map(|&(first, _)| first) + } + + fn leaf(&mut self, _node: &Arc>, first: usize, _end: usize) -> Arc> { + let (at, leaf) = self.0.next().expect("peeked just before"); + debug_assert_eq!(at, first, "a replacement leaf sits at its own position"); + leaf + } +} + /// The root at height `height` of the subtrees in `roots`, which sit at height /// `below`, left to right, with zero subtrees after them. /// @@ -380,7 +443,7 @@ fn fold_roots(mut layer: Vec, below: usize, height: usize) -> Hash256 { } /// A composite element's own root, as a leaf keeps it. -fn element_root(value: &T) -> Hash256 { +pub(crate) fn element_root(value: &T) -> Hash256 { HashTreeRoot::hash_tree_root(value, &Sha2Hasher) } diff --git a/crates/common/ssz-tree/src/vector.rs b/crates/common/ssz-tree/src/vector.rs index 0ec13690..45a1d023 100644 --- a/crates/common/ssz-tree/src/vector.rs +++ b/crates/common/ssz-tree/src/vector.rs @@ -7,6 +7,7 @@ use libssz::{DecodeError, SszDecode, SszEncode}; use libssz_merkle::{HashTreeRoot, Sha256Hasher}; use libssz_types::TypeError; +use crate::cursor::{ElemCow, IterCow}; use crate::interface::Interface; use crate::iter::Iter; use crate::update_map::{UpdateMap, VecMap}; @@ -48,6 +49,26 @@ impl> Vector { self.interface.get_mut(index) } + /// An in-order pass that can rewrite any element, for updating most of the + /// vector at once. Applies pending writes first; the pass's writes reach + /// the tree when it is dropped. See [`List::iter_cow`](crate::List::iter_cow). + pub fn iter_cow(&mut self) -> IterCow<'_, T> { + self.interface.iter_cow() + } + + /// Runs `f` on every element in order, stopping at the first error and + /// keeping the writes made before it. + pub fn try_update_each( + &mut self, + mut f: impl FnMut(&mut ElemCow<'_, T>) -> Result<(), E>, + ) -> Result<(), E> { + let mut pass = self.iter_cow(); + while let Some(mut element) = pass.next_cow() { + f(&mut element)?; + } + Ok(()) + } + /// The elements in order, pending writes included. pub fn iter(&self) -> Iter<'_, T, U> { self.interface.iter_from(0) diff --git a/crates/common/ssz-tree/tests/model.rs b/crates/common/ssz-tree/tests/model.rs index e09a5f1b..bdc6c857 100644 --- a/crates/common/ssz-tree/tests/model.rs +++ b/crates/common/ssz-tree/tests/model.rs @@ -521,3 +521,288 @@ proptest! { rebase_shrunk::>(orig_values, extra, hash_orig_first, hash_base_first)?; } } + +// ── Write cursor: random sweeps match a model, and clones are untouched ── + +/// What a sweep does to one element, picked by `action % 5`. +type Decision = (u8, T); + +fn decisions( + value: impl Strategy + Clone + 'static, +) -> impl Strategy>> { + vec((any::(), value), 1..9) +} + +/// Runs one sweep over `list`, mirroring every write in `model`, and stops +/// after `stop_at % (len + 1)` elements. Decision per element: skip; `set` a +/// new value; `set` the same value; `make_mut` and write a new value; +/// `make_mut` and write the same value. +fn sweep( + list: &mut List, + model: &mut [T], + stop_at: usize, + decisions: &[Decision], + noop_only: bool, +) -> Result<(), TestCaseError> +where + T: Value + Debug, + U: UpdateMap, +{ + let stop = stop_at % (model.len() + 1); + let mut pass = list.iter_cow(); + let mut seen = 0; + while let Some(mut element) = pass.next_cow() { + if seen == stop { + break; + } + let index = element.index(); + prop_assert_eq!(index, seen); + prop_assert_eq!(&*element, &model[index]); + let (action, value) = &decisions[index % decisions.len()]; + let action = if noop_only { + [0, 2, 4][*action as usize % 3] + } else { + *action % 5 + }; + match action { + 0 => {} + 1 => { + element.set(value.clone()); + model[index] = value.clone(); + } + 2 => { + let same = (*element).clone(); + element.set(same); + } + 3 => { + *element.make_mut() = value.clone(); + model[index] = value.clone(); + } + _ => { + let same = model[index].clone(); + *element.make_mut() = same; + } + } + seen += 1; + } + drop(pass); + Ok(()) +} + +fn run_sweeps( + initial: Vec, + ops: Vec>, + sweeps: Vec<(usize, Vec>)>, + hash_first: bool, +) -> Result<(), TestCaseError> +where + T: Value + Debug, + U: UpdateMap, +{ + let mut list = List::::try_from(initial.clone()).unwrap(); + let mut model = initial; + // Pending pushes and writes exercise the apply-first path. + apply_ops(&mut list, &mut model, ops); + if hash_first { + root(&list); + } + for (stop_at, decisions) in sweeps { + let before = list.clone(); + let before_model = model.clone(); + sweep(&mut list, &mut model, stop_at, &decisions, false)?; + check_list_matches(&list, &model)?; + // The clone taken before the sweep did not change. + check_list_matches(&before, &before_model)?; + if hash_first { + root(&list); + } + } + Ok(()) +} + +fn run_noop_sweep( + initial: Vec, + decisions: Vec>, + hash_first: bool, +) -> Result<(), TestCaseError> +where + T: Value + Debug, + U: UpdateMap, +{ + let mut list = List::::try_from(initial.clone()).unwrap(); + if hash_first { + root(&list); + } + let before = list.clone(); + let mut model = initial; + sweep(&mut list, &mut model, usize::MAX - 1, &decisions, true)?; + prop_assert!(list.ptr_eq(&before)); + check_list_matches(&list, &model) +} + +fn sweeps( + value: impl Strategy + Clone + 'static, +) -> impl Strategy>)>> { + vec((any::(), decisions(value)), 1..4) +} + +proptest! { + #[test] + fn u64_sweeps( + initial in vec(any::(), 0..40), + ops in ops(any::()), + sweeps in sweeps(any::()), + hash_first in any::(), + ) { + run_sweeps::>(initial, ops, sweeps, hash_first)?; + } + + #[test] + fn u64_sweeps_spanning_leaves( + initial in vec(any::(), 0..1600), + ops in ops(any::()), + sweeps in sweeps(any::()), + hash_first in any::(), + ) { + run_sweeps::>(initial, ops, sweeps, hash_first)?; + } + + #[test] + fn u64_sweeps_at_the_registry_limit( + initial in vec(any::(), 0..1100), + ops in ops(any::()), + sweeps in sweeps(any::()), + hash_first in any::(), + ) { + run_sweeps::>(initial, ops, sweeps, hash_first)?; + } + + /// A leaf holds 4096 u8s, the widest bitset. + #[test] + fn u8_sweeps_with_wide_leaves( + initial in vec(any::(), 0..4500), + sweeps in sweeps(any::()), + hash_first in any::(), + ) { + run_sweeps::>(initial, Vec::new(), sweeps, hash_first)?; + } + + #[test] + fn root_sweeps( + initial in vec(any::<[u8; 32]>(), 0..15), + ops in ops(any::<[u8; 32]>()), + sweeps in sweeps(any::<[u8; 32]>()), + hash_first in any::(), + ) { + run_sweeps::<[u8; 32], 20, BTreeMap>(initial, ops, sweeps, hash_first)?; + } + + #[test] + fn item_sweeps_spanning_leaves( + initial in vec(item(), 0..400), + ops in ops(item()), + sweeps in sweeps(item()), + hash_first in any::(), + ) { + run_sweeps::>(initial, ops, sweeps, hash_first)?; + } + + #[test] + fn blob_sweeps( + initial in vec(blob(), 0..6), + ops in ops(blob()), + sweeps in sweeps(blob()), + hash_first in any::(), + ) { + run_sweeps::>(initial, ops, sweeps, hash_first)?; + } + + #[test] + fn u64_noop_sweep_keeps_the_tree( + initial in vec(any::(), 0..1600), + decisions in decisions(any::()), + hash_first in any::(), + ) { + run_noop_sweep::>(initial, decisions, hash_first)?; + } + + #[test] + fn item_noop_sweep_keeps_the_tree( + initial in vec(item(), 0..400), + decisions in decisions(item()), + hash_first in any::(), + ) { + run_noop_sweep::>(initial, decisions, hash_first)?; + } + + #[test] + fn u64_vector_sweeps( + initial in vec(any::(), 1500), + decisions in decisions(any::()), + stop_at in any::(), + hash_first in any::(), + ) { + let mut vector = Vector::>::try_from(initial.clone()).unwrap(); + if hash_first { + root(&vector); + } + let mut model = initial; + let stop = stop_at % 1501; + let mut index = 0; + let mut pass = vector.iter_cow(); + while let Some(mut element) = pass.next_cow() { + if index == stop { + break; + } + let (action, value) = &decisions[index % decisions.len()]; + if action % 2 == 0 { + element.set(*value); + model[index] = *value; + } + index += 1; + } + drop(pass); + let reference = SszVector::::try_from(model.clone()).unwrap(); + prop_assert_eq!(vector.to_vec(), model); + prop_assert_eq!(root(&vector), root(&reference)); + } + + /// A `u64` list and an `Item` list swept in one loop, each writing from the + /// other's element. + #[test] + fn lockstep_sweep_over_two_lists( + ids in vec(any::(), 0..700), + data in any::<[u8; 32]>(), + hash_first in any::(), + ) { + let items: Vec = ids.iter().map(|&id| Item { id, data }).collect(); + let mut numbers = List::>::try_from(ids.clone()).unwrap(); + let mut list = List::>::try_from(items.clone()).unwrap(); + if hash_first { + root(&numbers); + root(&list); + } + let mut numbers_model = ids; + let mut items_model = items; + { + let mut numbers_pass = numbers.iter_cow(); + let mut items_pass = list.iter_cow(); + while let (Some(mut number), Some(mut item)) = + (numbers_pass.next_cow(), items_pass.next_cow()) + { + let index = number.index(); + prop_assert_eq!(index, item.index()); + if index % 3 == 0 { + number.set(item.id.wrapping_add(1)); + numbers_model[index] = items_model[index].id.wrapping_add(1); + } + if index % 2 == 0 { + item.make_mut().id = *number; + items_model[index].id = numbers_model[index]; + } + } + } + check_list_matches(&numbers, &numbers_model)?; + check_list_matches(&list, &items_model)?; + } +} diff --git a/crates/common/types/src/beacon/containers/mod.rs b/crates/common/types/src/beacon/containers/mod.rs index 40a240d9..a04149cc 100644 --- a/crates/common/types/src/beacon/containers/mod.rs +++ b/crates/common/types/src/beacon/containers/mod.rs @@ -395,6 +395,13 @@ impl BeaconState { } } +/// Disjoint mutable borrows of a state's validator registry and balances, from +/// [`BeaconState::registry_mut`]. +pub struct RegistryMut<'a> { + pub validators: &'a mut Validators, + pub balances: &'a mut Balances, +} + /// Generates read and write accessors for state fields that every fork shares. /// /// The `copy` and `reference` lists are this crate's statement of which state @@ -592,6 +599,20 @@ impl BeaconState { .ok_or(Error::UnknownValidator(index)) } + /// The validator registry and the balances, both mutably at once. + /// + /// `validators_mut` and `balances_mut` each borrow the whole state, so a pass + /// that reads one list while writing the other (an in-order write cursor + /// over one list zipped with an iterator over the other) needs this + /// accessor, the same way [`Self::altair_validator_lists_mut`] splits the + /// per-validator lists. + pub fn registry_mut(&mut self) -> RegistryMut<'_> { + dispatch_state!(self, "BeaconState::registry_mut", |state| RegistryMut { + validators: &mut state.validators, + balances: &mut state.balances, + }) + } + /// The randao mix for `epoch`, which the specification indexes modulo the /// vector length so the vector acts as a ring buffer. pub fn randao_mix(&self, epoch: Epoch) -> Bytes32 { diff --git a/crates/common/types/tests/tree_bench.rs b/crates/common/types/tests/tree_bench.rs index e560b3d3..52937e72 100644 --- a/crates/common/types/tests/tree_bench.rs +++ b/crates/common/types/tests/tree_bench.rs @@ -203,6 +203,127 @@ fn tree_bench() { root(&balances) }); + // Where `epoch_sweep_tree`'s time goes: filling the pending-write map, + // folding it into the tree, hashing the result. + { + let mut balances = tree_balances.clone(); + let start = Instant::now(); + for index in 0..VALIDATOR_COUNT { + balances[index] += 1; + } + let write = start.elapsed(); + let start = Instant::now(); + balances.apply_updates(); + let apply = start.elapsed(); + let start = Instant::now(); + black_box(root(&balances)); + println!( + "tree_bench epoch_sweep_tree_phases write {:.1} ms apply {:.1} ms hash {:.1} ms", + write.as_secs_f64() * 1e3, + apply.as_secs_f64() * 1e3, + start.elapsed().as_secs_f64() * 1e3 + ); + } + + // The same sweep through the write cursor: no pending-write map, no sorted + // copy of it. + let sweep_root = time("epoch_sweep_cursor_tree", || { + let mut balances = tree_balances.clone(); + balances + .try_update_each(|balance| { + balance.set(**balance + 1); + Ok::<_, ()>(()) + }) + .unwrap(); + root(&balances) + }); + let mut expected_sweep = tree_balances.clone(); + for index in 0..VALIDATOR_COUNT { + expected_sweep[index] += 1; + } + expected_sweep.apply_updates(); + assert_eq!( + sweep_root, + root(&expected_sweep), + "the cursor must write the same" + ); + drop(expected_sweep); + + // Nothing changes, so every leaf and its hash are kept. + time("epoch_sweep_noop_cursor_tree", || { + let mut balances = tree_balances.clone(); + balances + .try_update_each(|balance| { + let same = **balance; + balance.set(same); + Ok::<_, ()>(()) + }) + .unwrap(); + assert!(balances.ptr_eq(&tree_balances)); + root(&balances) + }); + + // The rebuild alternative: copy out, write, build a new list, share what is + // equal with the old one. + time("epoch_sweep_rebuild_tree", || { + let mut values = tree_balances.to_vec(); + for value in &mut values { + *value += 1; + } + let mut balances = Balances::try_from(values).unwrap(); + balances.rebase_on(&tree_balances); + root(&balances) + }); + + // The rewards step: four (reward, penalty) pairs per validator, each + // applied as a saturating add then subtract, in order. + let deltas: Vec<(Vec, Vec)> = (0..4u64) + .map(|pair| { + let rewards = (0..VALIDATOR_COUNT as u64) + .map(|i| (i + pair) % 7) + .collect(); + let penalties = (0..VALIDATOR_COUNT as u64) + .map(|i| (i * 3 + pair) % 5) + .collect(); + (rewards, penalties) + }) + .collect(); + let get_mut_root = time("rewards_apply_4_pairs_get_mut", || { + let mut balances = tree_balances.clone(); + for (rewards, penalties) in &deltas { + for index in 0..VALIDATOR_COUNT { + let balance = balances.get_mut(index).unwrap(); + *balance = balance.saturating_add(rewards[index]); + let balance = balances.get_mut(index).unwrap(); + *balance = balance.saturating_sub(penalties[index]); + } + } + balances.apply_updates(); + root(&balances) + }); + let cursor_root = time("rewards_apply_4_pairs_cursor", || { + let mut balances = tree_balances.clone(); + balances + .try_update_each(|balance| { + let index = balance.index(); + let mut value = **balance; + for (rewards, penalties) in &deltas { + value = value + .saturating_add(rewards[index]) + .saturating_sub(penalties[index]); + } + balance.set(value); + Ok::<_, ()>(()) + }) + .unwrap(); + root(&balances) + }); + assert_eq!( + get_mut_root, cursor_root, + "the cursor must apply the same deltas" + ); + drop(deltas); + // Memory held by a chain of derived states, per state. let before = live_mib(); let mut tree_states = Vec::with_capacity(TREE_STATES); @@ -312,4 +433,58 @@ fn tree_bench() { .wrapping_add(validator.effective_balance) }) }); + + // `process_effective_balance_updates` with 0.1% of the validators changing: + // the old indexed read plus buffered writes, against the cursor over the + // registry zipped with the balances. + let indexed_root = time("eb_updates_indexed", || { + let mut validators = tree_validators.clone(); + let mut updates = Vec::new(); + for (index, validator) in validators.iter().enumerate() { + let balance = tree_balances[index]; + if index % 1000 == 0 && balance != validator.effective_balance { + updates.push((index, validator.effective_balance + 1)); + } + } + for (index, effective) in updates { + validators[index].effective_balance = effective; + } + validators.apply_updates(); + root(&validators) + }); + let cursor_root = time("eb_updates_cursor", || { + let mut validators = tree_validators.clone(); + let mut balances = tree_balances.iter(); + let mut pass = validators.iter_cow(); + while let Some(mut validator) = pass.next_cow() { + let balance = *balances.next().unwrap(); + if validator.index() % 1000 == 0 && balance != validator.effective_balance { + let effective = validator.effective_balance + 1; + validator.make_mut().effective_balance = effective; + } + } + drop(pass); + root(&validators) + }); + assert_eq!(indexed_root, cursor_root, "the cursor must write the same"); + + // Validators, balances and a third list swept in one loop, as a single + // epoch pass would: one cursor each, only the third one written. + time("lockstep_three_lists_cursor", || { + let (mut validators, mut balances) = (tree_validators.clone(), tree_balances.clone()); + let mut scores = tree_balances.clone(); + { + let mut validators = validators.iter_cow(); + let mut balances = balances.iter_cow(); + let mut scores = scores.iter_cow(); + while let (Some(validator), Some(balance), Some(mut score)) = ( + validators.next_cow(), + balances.next_cow(), + scores.next_cow(), + ) { + score.set(validator.effective_balance.wrapping_add(*balance)); + } + } + (root(&validators), root(&balances), root(&scores)) + }); } diff --git a/docs/beacon_stf.md b/docs/beacon_stf.md index 17ebe686..a7b0a579 100644 --- a/docs/beacon_stf.md +++ b/docs/beacon_stf.md @@ -279,6 +279,21 @@ lists lighthouse keeps its state in. writes pending is still correct but computed on a throwaway copy, and the store flushes a state before caching it, since a shared `Arc` cannot be flushed later. +- **A pass that writes most of a list uses the write cursor instead.** + `List::iter_cow` (and `try_update_each`, its closure form) walks the list in + order and copies a leaf only when a write in it changes something, comparing + against the original at the leaf's end and keeping the original leaf, its + `Arc` and its hash if nothing differs. Nothing goes through the buffer, so a + full sweep avoids a map entry and a sorted copy per element. The changed leaves are swapped into the tree + when the cursor is dropped, which rebuilds only the paths to them; a changed + composite element's root is computed then, on the state-transition thread. Writes pending + when the pass starts are applied first. `BeaconState::registry_mut` hands out + `validators` and `balances` as disjoint borrows, so a pass can write one while + reading the other. The balance application of `process_rewards_and_penalties` + and both `process_effective_balance_updates` run this way: they move + together, since once the rewards write bypasses the buffer, the effective + balance pass would otherwise pay a tree descent per `balances()[i]` instead of a + buffer hit. The access pattern matters. `state.validator(i)` and `balances()[i]` are tree descents, cheap next to a hash but far from an array index, and they add up