diff --git a/CLAUDE.md b/CLAUDE.md index ebef603a9..a151a49d9 100644 --- a/CLAUDE.md +++ b/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 diff --git a/crates/blockchain/state_transition/src/beacon/helpers/mutators.rs b/crates/blockchain/state_transition/src/beacon/helpers/mutators.rs index 077591fc4..c9ae7ff45 100644 --- a/crates/blockchain/state_transition/src/beacon/helpers/mutators.rs +++ b/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, Vec)], +) -> 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 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 6aadf09a5..7ac49fee4 100644 --- a/crates/blockchain/state_transition/src/beacon/stf/epoch/altair.rs +++ b/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. 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 6836ad495..672d78663 100644 --- a/crates/blockchain/state_transition/src/beacon/stf/epoch/electra.rs +++ b/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` 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 86c853e75..6d59a5394 100644 --- a/crates/blockchain/state_transition/src/beacon/stf/epoch/mod.rs +++ b/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 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 22bebbe3b..4f5515bc1 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.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)] diff --git a/crates/common/ssz-tree/src/cursor.rs b/crates/common/ssz-tree/src/cursor.rs new file mode 100644 index 000000000..5ee3f27fd --- /dev/null +++ b/crates/common/ssz-tree/src/cursor.rs @@ -0,0 +1,584 @@ +//! 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>)>, + /// Added to every index handed out: where this tree starts in a list made + /// of several (a progressive list's subtrees). + base: usize, +} + +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(), + base: 0, + } + } + + /// Reports indices offset by `base`, for a tree that is not the start of + /// its list. + pub(crate) fn with_base(mut self, base: usize) -> Self { + self.base = base; + self + } + + /// The next element, to read or to write, or `None` after the last. + pub fn next_cow(&mut self) -> Option> { + if self.advance() { + Some(self.take()) + } else { + None + } + } + + /// Positions the pass on the next element, loading the next leaf when the + /// current one is spent. `false` once the data ends. + /// + /// Split from [`IterCow::take`] so a caller chaining several passes can + /// ask whether one has an element without holding a borrow of it. + pub(crate) fn advance(&mut self) -> bool { + while self.state.next >= self.state.len { + self.finish_leaf(); + let Some(leaf) = self.next_leaf() else { + return false; + }; + let first = self.next_first; + self.next_first += crate::packing_factor::() << crate::leaf_height::(self.depth); + self.state.load(leaf, first); + } + true + } + + /// The element [`IterCow::advance`] positioned on. + pub(crate) fn take(&mut self) -> ElemCow<'_, T> { + let offset = self.state.next; + self.state.next += 1; + let index = self.base + self.state.first + offset; + ElemCow { + state: &mut self.state, + offset, + index, + } + } + + 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); + } + + #[test] + fn a_progressive_list_pass_crosses_subtrees_with_global_indices() { + use crate::ProgressiveList; + // 3000 u64s span several subtrees (4, 16, 64, ... chunks). + let values: Vec = (0..3000).collect(); + let mut list = ProgressiveList::::from(values.clone()); + *list.get_mut(7).unwrap() = 70; + list.push(3000).unwrap(); + list.try_update_each(|e| -> Result<(), ()> { + let index = e.index(); + assert_eq!(**e, if index == 7 { 70 } else { index as u64 }); + if index % 3 == 0 { + e.set(index as u64 + 1); + } + Ok(()) + }) + .unwrap(); + let expected: Vec = (0..=3000u64) + .map(|i| match i { + 7 => 70, + i if i % 3 == 0 => i + 1, + i => i, + }) + .collect(); + assert_eq!(list.to_vec(), expected); + let reference = libssz_types::ProgressiveList::from(expected); + assert_eq!(list_root(&list), list_root(&reference)); + } +} diff --git a/crates/common/ssz-tree/src/interface.rs b/crates/common/ssz-tree/src/interface.rs index 671f6c60c..19f3b8fff 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 34b8efae2..bb3866152 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 @@ -61,6 +71,7 @@ //! right, so writes, hash caching and rebasing all work per subtree exactly //! as they do for [`List`]. +mod cursor; mod interface; mod iter; mod list; @@ -70,9 +81,10 @@ mod tree; mod update_map; mod vector; +pub use cursor::{ElemCow, IterCow}; pub use iter::Iter; pub use list::List; -pub use progressive_list::{ProgressiveIter, ProgressiveList}; +pub use progressive_list::{ProgressiveIter, ProgressiveIterCow, ProgressiveList}; pub use update_map::{UpdateMap, VecMap}; pub use vector::Vector; diff --git a/crates/common/ssz-tree/src/list.rs b/crates/common/ssz-tree/src/list.rs index f0f98b6f1..62efc6ab8 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/progressive_list.rs b/crates/common/ssz-tree/src/progressive_list.rs index 081dd1206..861b5b187 100644 --- a/crates/common/ssz-tree/src/progressive_list.rs +++ b/crates/common/ssz-tree/src/progressive_list.rs @@ -9,6 +9,7 @@ use libssz::{BYTES_PER_LENGTH_OFFSET, DecodeError, SszDecode, SszEncode}; use libssz_merkle::{HashTreeRoot, Sha2Hasher, Sha256Hasher, hash_nodes, 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}; @@ -132,6 +133,32 @@ impl> ProgressiveList { } } + /// An in-order pass that can rewrite any element, subtree after subtree. + /// See [`List::iter_cow`](crate::List::iter_cow): pending writes are + /// applied first, and each subtree's changed leaves reach it when the + /// pass leaves that subtree or is dropped. + pub fn iter_cow(&mut self) -> ProgressiveIterCow<'_, T, U> { + self.apply_updates(); + ProgressiveIterCow { + subtrees: self.subtrees.iter_mut(), + current: None, + base: 0, + } + } + + /// 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, copied out. pub fn to_vec(&self) -> Vec { self.iter().cloned().collect() @@ -177,6 +204,32 @@ impl> ProgressiveList { } } +/// An in-order rewriting pass over a [`ProgressiveList`], made by +/// [`ProgressiveList::iter_cow`]: one [`IterCow`] per subtree, in turn. +pub struct ProgressiveIterCow<'a, T: Value, U> { + subtrees: std::slice::IterMut<'a, Interface>, + current: Option>, + /// Elements in the subtrees already started. + base: usize, +} + +impl> ProgressiveIterCow<'_, T, U> { + /// The next element, to read or to write, or `None` after the last. + pub fn next_cow(&mut self) -> Option> { + loop { + if self.current.as_mut().is_some_and(IterCow::advance) { + return self.current.as_mut().map(IterCow::take); + } + // Dropping the spent pass swaps its changed leaves into its subtree. + self.current = None; + let subtree = self.subtrees.next()?; + let len = subtree.len(); + self.current = Some(subtree.iter_cow().with_base(self.base)); + self.base += len; + } + } +} + impl> Default for ProgressiveList { fn default() -> Self { Self::empty() diff --git a/crates/common/ssz-tree/src/tree.rs b/crates/common/ssz-tree/src/tree.rs index 2eb8c70d3..9339e722c 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 0ec13690c..45a1d0238 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 f41652401..a56db8597 100644 --- a/crates/common/ssz-tree/tests/model.rs +++ b/crates/common/ssz-tree/tests/model.rs @@ -994,3 +994,288 @@ fn progressive_rebase_ignores_a_pending_push_into_a_new_subtree() { orig.rebase_on(&base); assert!(orig.ptr_eq(&base)); } + +// ── 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 367da3c75..201ed27c9 100644 --- a/crates/common/types/src/beacon/containers/mod.rs +++ b/crates/common/types/src/beacon/containers/mod.rs @@ -821,6 +821,56 @@ impl BeaconState { } } + /// Runs `f` over every balance in order through the list's write cursor + /// (`ethlambda_ssz_tree::List::iter_cow`, or the progressive list's + /// equivalent from gloas on), stopping at the first error and keeping the + /// writes made before it. + /// + /// For a pass that rewrites most balances: nothing is buffered per element, + /// and a leaf whose balances all come out unchanged keeps its hash. + pub fn try_update_balances( + &mut self, + mut f: impl FnMut(&mut ethlambda_ssz_tree::ElemCow<'_, Gwei>) -> core::result::Result<(), E>, + ) -> core::result::Result<(), E> { + dispatch_state!(self, "BeaconState::try_update_balances", |state| state + .balances + .try_update_each(&mut f)) + } + + /// Runs `f` over every validator in order through the registry's write + /// cursor, with that validator's balance read in step. + /// + /// `validators` and `balances` are separate lists of the same state, so a + /// pass that rewrites one while reading the other cannot go through + /// [`Self::validator_mut`] and [`Self::balance`] one call at a time: this + /// splits the borrow. It stops at the first error, keeping the writes made + /// before it, and also stops (without error) when the balances run out + /// before the validators do: the caller decides whether a short balances + /// list is an error. + pub fn try_update_validators_with_balances( + &mut self, + mut f: impl FnMut( + &mut ethlambda_ssz_tree::ElemCow<'_, Validator>, + Gwei, + ) -> core::result::Result<(), E>, + ) -> core::result::Result<(), E> { + dispatch_state!( + self, + "BeaconState::try_update_validators_with_balances", + |state| { + let mut balances = state.balances.iter(); + let mut pass = state.validators.iter_cow(); + while let Some(mut validator) = pass.next_cow() { + let Some(&balance) = balances.next() else { + break; + }; + f(&mut validator, balance)?; + } + Ok(()) + } + ) + } + /// 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 8f1421325..41d59e10e 100644 --- a/crates/common/types/tests/tree_bench.rs +++ b/crates/common/types/tests/tree_bench.rs @@ -205,6 +205,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); @@ -314,6 +435,60 @@ 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)) + }); } /// Gloas turns the validator registry into a `ProgressiveList` (EIP-7916). diff --git a/docs/beacon_stf.md b/docs/beacon_stf.md index 313da7e70..90783ac1b 100644 --- a/docs/beacon_stf.md +++ b/docs/beacon_stf.md @@ -289,6 +289,21 @@ below): persistent Merkle trees in the shape of the SSZ one, modeled on the 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 `state.balance(i)` are tree descents, cheap next to a hash but far from an array index, and they add up