From d9079526927dda9942b4b8dc50066624454e1adc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Gr=C3=BCner?= <47506558+MegaRedHand@users.noreply.github.com> Date: Fri, 2 Oct 2026 11:50:17 -0300 Subject: [PATCH 1/2] perf(ssz-tree): add an in-order copy-on-write write cursor Rewriting most of a list through get_mut buffers an entry per element, sorts the buffer into a second Vec, and then rebuilds every touched leaf, recomputing the root of every written element even when the value did not change. iter_cow / try_update_each walk the leaves in order instead. A leaf is copied on the first write that reaches it, compared with the original at the leaf's end (only the offsets make_mut handed out), and kept as it was (Arc and cached hash) if nothing differs. Changed leaves carry over the element roots of unchanged elements and compute the changed ones. The changed leaves are swapped in on drop with one descent, so an early return or a panic leaves a consistent list. Pending writes are applied before the pass starts. The with_updated_leaves descent is generalized over a leaf source so the element-update path and the whole-leaf replacement share it. Covered by unit tests (changed-path sharing, carried roots, early return, catch_unwind, registry depth) and model-based property tests against SszList. --- crates/common/ssz-tree/src/cursor.rs | 526 ++++++++++++++++++++++++ crates/common/ssz-tree/src/interface.rs | 11 + crates/common/ssz-tree/src/lib.rs | 12 + crates/common/ssz-tree/src/list.rs | 22 + crates/common/ssz-tree/src/tree.rs | 81 +++- crates/common/ssz-tree/src/vector.rs | 21 + crates/common/ssz-tree/tests/model.rs | 285 +++++++++++++ 7 files changed, 949 insertions(+), 9 deletions(-) create mode 100644 crates/common/ssz-tree/src/cursor.rs 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)?; + } +} From 14efa29e687c551dba5e865f1c434cf92086e958 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Gr=C3=BCner?= <47506558+MegaRedHand@users.noreply.github.com> Date: Fri, 2 Oct 2026 11:50:30 -0300 Subject: [PATCH 2/2] perf(beacon): apply epoch balance and effective-balance writes through the cursor Both rewards-and-penalties passes wrote every balance through get_mut, up to eight calls per validator on altair and later, each a map entry in the dense pending-write buffer. They now apply every (rewards, penalties) pair per index in one cursor pass, keeping the spec's per-pair saturating_add then saturating_sub order (no netting). Both process_effective_balance_updates variants read balances()[index] per validator, which was a buffer hit only because the rewards pass had just filled the buffer. With the rewards write bypassing the buffer those reads would become tree descents, so they move in the same change: the registry's cursor runs zipped with a balances iterator, through the new BeaconState::registry_mut disjoint borrow, and copies a leaf only for a validator whose effective balance actually changes. Adds the tree_bench rows comparing the old and new write paths, and the docs for the cursor. --- CLAUDE.md | 5 +- .../src/beacon/helpers/mutators.rs | 42 +++++ .../src/beacon/stf/epoch/altair.rs | 24 +-- .../src/beacon/stf/epoch/electra.rs | 37 ++-- .../src/beacon/stf/epoch/mod.rs | 37 ++-- .../src/beacon/stf/epoch/rewards.rs | 20 +- .../common/types/src/beacon/containers/mod.rs | 21 +++ crates/common/types/tests/tree_bench.rs | 175 ++++++++++++++++++ docs/beacon_stf.md | 15 ++ 9 files changed, 315 insertions(+), 61 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 74d3544b..d93b7071 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -584,7 +584,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/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