From f1272429db989d4ce4990fe5bfffcc0cee7cdb14 Mon Sep 17 00:00:00 2001 From: Derek Cofausper <256792747+decofe@users.noreply.github.com> Date: Wed, 4 Mar 2026 05:59:14 -0800 Subject: [PATCH] =?UTF-8?q?chore(trie):=20proof=5Fv2=20cleanup=20=E2=80=94?= =?UTF-8?q?=20use=20Nibbles/TrieMask=20builtins=20(#22769)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Brian Picciano <933154+mediocregopher@users.noreply.github.com> Co-authored-by: Brian Picciano --- .changelog/nice-bears-cook.md | 5 ++ crates/trie/trie/src/proof_v2/mod.rs | 74 +++---------------------- crates/trie/trie/src/proof_v2/target.rs | 3 +- 3 files changed, 15 insertions(+), 67 deletions(-) create mode 100644 .changelog/nice-bears-cook.md diff --git a/.changelog/nice-bears-cook.md b/.changelog/nice-bears-cook.md new file mode 100644 index 0000000000..2c9e30bbbd --- /dev/null +++ b/.changelog/nice-bears-cook.md @@ -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`. diff --git a/crates/trie/trie/src/proof_v2/mod.rs b/crates/trie/trie/src/proof_v2/mod.rs index a270c43296..9b050cad3f 100644 --- a/crates/trie/trie/src/proof_v2/mod.rs +++ b/crates/trie/trie/src/proof_v2/mod.rs @@ -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)), } -/// 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 { - 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)> = 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}; diff --git a/crates/trie/trie/src/proof_v2/target.rs b/crates/trie/trie/src/proof_v2/target.rs index f5a0f6d679..14be86c7ac 100644 --- a/crates/trie/trie/src/proof_v2/target.rs +++ b/crates/trie/trie/src/proof_v2/target.rs @@ -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 { - 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