chore(trie): proof_v2 cleanup — use Nibbles/TrieMask builtins (#22769)

Co-authored-by: Brian Picciano <933154+mediocregopher@users.noreply.github.com>
Co-authored-by: Brian Picciano <me@mediocregopher.com>
This commit is contained in:
Derek Cofausper
2026-03-04 05:59:14 -08:00
committed by GitHub
parent ad96bc4649
commit f1272429db
3 changed files with 15 additions and 67 deletions

View File

@@ -0,0 +1,5 @@
---
reth-trie: patch
---
Removed the local `increment_and_strip_trailing_zeros` function and `PATH_ALL_ZEROS` static in `proof_v2`, replacing them with the equivalent `Nibbles::next_without_prefix` and `Nibbles::is_zeroes` builtins. Also replaced manual `.get()` calls on `state_mask`/`hash_mask` with direct field access and switched to `Nibbles::unpack_array` over the unsafe `unpack_unchecked`.

View File

@@ -37,17 +37,6 @@ static TRACE_TARGET: &str = "trie::proof_v2";
/// Number of bytes to pre-allocate for [`ProofCalculator`]'s `rlp_encode_buf` field.
const RLP_ENCODE_BUF_SIZE: usize = 1024;
/// A [`Nibbles`] which contains 64 zero nibbles.
static PATH_ALL_ZEROS: Nibbles = {
let mut path = Nibbles::new();
let mut i = 0;
while i < 64 {
path.push_unchecked(0);
i += 1;
}
path
};
/// A proof calculator that generates merkle proofs using only leaf data.
///
/// The calculator:
@@ -675,8 +664,7 @@ where
// leaf value can begin ASAP.
let mut map_hashed_cursor_entry = |(key_b256, val): (B256, _)| {
debug_assert_eq!(key_b256.len(), 32);
// SAFETY: key is a B256 and so is exactly 32-bytes.
let key = unsafe { Nibbles::unpack_unchecked(key_b256.as_slice()) };
let key = Nibbles::unpack_array(key_b256.as_ref());
let val = value_encoder.deferred_encoder(key_b256, val);
(key, val)
};
@@ -883,7 +871,7 @@ where
// If the next cached branch's path is all zeros then we can skip this catch-up step,
// because there cannot be any keys prior to that range.
let cached_path = &cached.0;
if uncalculated_lower_bound < cached_path && !PATH_ALL_ZEROS.starts_with(cached_path) {
if uncalculated_lower_bound < cached_path && !cached_path.is_zeroes() {
let range = (*uncalculated_lower_bound, Some(*cached_path));
trace!(target: TRACE_TARGET, ?range, "Returning key range to calculate in order to catch up to cached branch");
@@ -987,8 +975,8 @@ where
let curr_branch =
self.branch_stack.last().expect("top of branch_stack corresponds to cached branch");
let cached_state_mask = cached_branch.state_mask.get();
let curr_state_mask = curr_branch.state_mask.get();
let cached_state_mask = cached_branch.state_mask;
let curr_state_mask = curr_branch.state_mask;
// Determine all child nibbles which are set in the cached branch but not the
// under-construction branch.
@@ -1001,7 +989,7 @@ where
// If there are no further children to construct for this branch then pop it off both
// stacks and loop using the parent branch.
if next_child_nibbles == 0 {
if next_child_nibbles.is_empty() {
trace!(
target: TRACE_TARGET,
path=?cached_path,
@@ -1017,7 +1005,7 @@ where
// The just-popped branch is completely processed; we know there can be no more keys
// with that prefix. Set the lower bound which can be returned from this method to
// be the next possible prefix, if any.
uncalculated_lower_bound = increment_and_strip_trailing_zeros(&cached_path);
uncalculated_lower_bound = cached_path.next_without_prefix();
continue
}
@@ -1047,7 +1035,7 @@ where
// we first need to calculate the mask of which cached hashes have already been
// used by this branch (if any). The number of set bits in that mask will be the
// index of the next hash in the array to use.
let curr_hashed_used_mask = cached_branch.hash_mask.get() & curr_state_mask;
let curr_hashed_used_mask = cached_branch.hash_mask & curr_state_mask;
let hash_idx = curr_hashed_used_mask.count_ones() as usize;
let hash = cached_branch.hashes[hash_idx];
@@ -1068,7 +1056,7 @@ where
// Update the `uncalculated_lower_bound` to indicate that the child whose bit
// was just set is completely processed.
uncalculated_lower_bound = increment_and_strip_trailing_zeros(&child_path);
uncalculated_lower_bound = child_path.next_without_prefix();
// Push the current cached branch back onto the stack before looping.
self.cached_branch_stack.push((cached_path, cached_branch));
@@ -1112,7 +1100,7 @@ where
// There is no cached data for the sub-trie at this child, we must recalculate the
// sub-trie root (this child) using the leaves. Return the range of keys based on the
// child path.
let child_path_upper = increment_and_strip_trailing_zeros(&child_path);
let child_path_upper = child_path.next_without_prefix();
trace!(
target: TRACE_TARGET,
lower=?child_path,
@@ -1617,26 +1605,6 @@ enum PopCachedBranchOutcome {
CalculateLeaves((Nibbles, Option<Nibbles>)),
}
/// Increments the nibbles and strips any trailing zeros.
///
/// This function wraps `Nibbles::increment` and when it returns a value with trailing zeros,
/// it strips those zeros using bit manipulation on the underlying U256.
fn increment_and_strip_trailing_zeros(nibbles: &Nibbles) -> Option<Nibbles> {
let mut result = nibbles.increment()?;
// If result is empty, just return it
if result.is_empty() {
return Some(result);
}
// Get access to the underlying U256 to detect trailing zeros
let uint_val = *result.as_mut_uint_unchecked();
let non_zero_prefix_len = 64 - (uint_val.trailing_zeros() / 4);
result.truncate(non_zero_prefix_len);
Some(result)
}
#[cfg(test)]
mod tests {
use super::*;
@@ -1989,30 +1957,6 @@ mod tests {
.expect("Proof generation failed");
}
#[test]
fn test_increment_and_strip_trailing_zeros() {
let test_cases: Vec<(Nibbles, Option<Nibbles>)> = vec![
// Basic increment without trailing zeros
(Nibbles::from_nibbles([0x1, 0x2, 0x3]), Some(Nibbles::from_nibbles([0x1, 0x2, 0x4]))),
// Increment with trailing zeros - should be stripped
(Nibbles::from_nibbles([0x0, 0x0, 0xF]), Some(Nibbles::from_nibbles([0x0, 0x1]))),
(Nibbles::from_nibbles([0x0, 0xF, 0xF]), Some(Nibbles::from_nibbles([0x1]))),
// Overflow case
(Nibbles::from_nibbles([0xF, 0xF, 0xF]), None),
// Empty nibbles
(Nibbles::new(), None),
// Single nibble
(Nibbles::from_nibbles([0x5]), Some(Nibbles::from_nibbles([0x6]))),
// All Fs except last - results in trailing zeros after increment
(Nibbles::from_nibbles([0xE, 0xF, 0xF]), Some(Nibbles::from_nibbles([0xF]))),
];
for (input, expected) in test_cases {
let result = increment_and_strip_trailing_zeros(&input);
assert_eq!(result, expected, "Failed for input: {:?}", input);
}
}
#[test]
fn test_failing_proptest_case_0() {
use alloy_primitives::{hex, map::B256Map};

View File

@@ -1,4 +1,3 @@
use crate::proof_v2::increment_and_strip_trailing_zeros;
use reth_trie_common::{Nibbles, ProofV2Target};
// A helper function for getting the largest prefix of the sub-trie which contains a particular
@@ -29,7 +28,7 @@ pub(crate) fn sub_trie_prefix(target: &ProofV2Target) -> Nibbles {
// A helper function which returns the first path following a sub-trie in lexicographical order.
#[inline]
fn sub_trie_upper_bound(sub_trie_prefix: &Nibbles) -> Option<Nibbles> {
increment_and_strip_trailing_zeros(sub_trie_prefix)
sub_trie_prefix.next_without_prefix()
}
/// Describes a set of targets which all apply to a single sub-trie, ie a section of the overall