feat(trie): Proof Rewrite: Support partial proofs (#20336)

Co-authored-by: YK <chiayongkang@hotmail.com>
This commit is contained in:
Brian Picciano
2025-12-23 13:42:07 +01:00
committed by GitHub
parent 9f2aea0494
commit b79c58d835
4 changed files with 1062 additions and 205 deletions

View File

@@ -161,10 +161,10 @@ fn bench_proof_algos(c: &mut Criterion) {
StorageProofCalculator::new_storage(trie_cursor, hashed_cursor);
b.iter_batched(
|| targets.clone(),
|targets| {
|| targets.iter().copied().map(Into::into).collect::<Vec<_>>(),
|mut targets| {
proof_calculator
.storage_proof(hashed_address, targets)
.storage_proof(hashed_address, &mut targets)
.expect("Proof generation failed");
},
BatchSize::SmallInput,

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,335 @@
use crate::proof_v2::increment_and_strip_trailing_zeros;
use alloy_primitives::B256;
use reth_trie_common::Nibbles;
/// Target describes a proof target. For every proof target given, the
/// [`crate::proof_v2::ProofCalculator`] will calculate and return all nodes whose path is a prefix
/// of the target's `key`.
#[derive(Debug, Copy, Clone)]
pub struct Target {
pub(crate) key: Nibbles,
pub(crate) min_len: u8,
}
impl Target {
/// Returns a new [`Target`] which matches all trie nodes whose path is a prefix of this key.
pub fn new(key: B256) -> Self {
// SAFETY: key is a B256 and so is exactly 32-bytes.
let key = unsafe { Nibbles::unpack_unchecked(key.as_slice()) };
Self { key, min_len: 0 }
}
/// Only match trie nodes whose path is at least this long.
///
/// # Panics
///
/// This method panics if `min_len` is greater than 64.
pub fn with_min_len(mut self, min_len: u8) -> Self {
debug_assert!(min_len <= 64);
self.min_len = min_len;
self
}
// A helper function for getting the largest prefix of the sub-trie which contains a particular
// target, based on its `min_len`.
//
// A target will only match nodes which share the target's prefix, where the target's prefix is
// the first `min_len` nibbles of its key. E.g. a target with `key` 0xabcd and `min_len` 2 will
// only match nodes with prefix 0xab.
//
// In general the target will only match within the sub-trie whose prefix is identical to the
// target's. However there is an exception:
//
// Given a trie with a node at 0xabc, there must be a branch at 0xab. A target with prefix 0xabc
// needs to match that node, but the branch at 0xab must be constructed order to know the node
// is at that path. Therefore the sub-trie prefix is the target prefix with a nibble truncated.
//
// For a target with an empty prefix (`min_len` of 0) we still use an empty sub-trie prefix;
// this will still construct the branch at the root node (if there is one). Targets with
// `min_len` of both 0 and 1 will therefore construct the root node, but only those with
// `min_len` of 0 will retain it.
#[inline]
fn sub_trie_prefix(&self) -> Nibbles {
let mut sub_trie_prefix = self.key;
sub_trie_prefix.truncate(self.min_len.saturating_sub(1) as usize);
sub_trie_prefix
}
}
impl From<B256> for Target {
fn from(key: B256) -> Self {
Self::new(key)
}
}
// 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)
}
/// Describes a set of targets which all apply to a single sub-trie, ie a section of the overall
/// trie whose nodes all share a prefix.
pub(crate) struct SubTrieTargets<'a> {
/// The prefix which all nodes in the sub-trie share. This is also the first node in the trie
/// in lexicographic order.
pub(crate) prefix: Nibbles,
/// The targets belonging to this sub-trie. These will be sorted by their `key` field,
/// lexicographically.
pub(crate) targets: &'a [Target],
/// Will be true if at least one target in the set has a zero `min_len`.
///
/// If this is true then `prefix.is_empty()`, though not necessarily vice-versa.
pub(crate) retain_root: bool,
}
impl<'a> SubTrieTargets<'a> {
// A helper function which returns the first path following a sub-trie in lexicographical order.
#[inline]
pub(crate) fn upper_bound(&self) -> Option<Nibbles> {
sub_trie_upper_bound(&self.prefix)
}
}
/// Given a set of [`Target`]s, returns an iterator over those same [`Target`]s chunked by the
/// sub-tries they apply to within the overall trie.
pub(crate) fn iter_sub_trie_targets<'a>(
targets: &'a mut [Target],
) -> impl Iterator<Item = SubTrieTargets<'a>> {
// First sort by the sub-trie prefix of each target, falling back to the `min_len` in cases
// where the sub-trie prefixes are equal (to differentiate targets which match the root node and
// those which don't).
targets.sort_unstable_by(|a, b| {
a.sub_trie_prefix().cmp(&b.sub_trie_prefix()).then_with(|| a.min_len.cmp(&b.min_len))
});
// We now chunk targets, such that each chunk contains all targets belonging to the same
// sub-trie. We are taking advantage of the following properties:
//
// - The first target in the chunk has the shortest sub-trie prefix (see previous sorting step).
//
// - The upper bound of the first target in the chunk's sub-trie will therefore be the upper
// bound of the whole chunk.
// - For example, given a chunk with sub-trie prefixes [0x2, 0x2f, 0x2fa], the upper bounds
// will be [0x3, 0x3, 0x2fb]. Note that no target could match a trie node with path equal
// to or greater than 0x3.
//
// - If a target's sub-trie's prefix does not lie within the bounds of the current chunk, then
// that target must be the first target of the next chunk, lying in a separate sub-trie.
// - Example: given sub-trie prefixes of [0x2, 0x2fa, 0x4c, 0x4ce, 0x4e], we would end up
// with the following chunks:
// - [0x2, 0x2fa] w/ upper bound 0x3
// - [0x4c 0x4ce] w/ upper bound 0x4d
// - [0x4e] w/ upper bound 0x4f
let mut upper_bound = targets.first().and_then(|t| sub_trie_upper_bound(&t.sub_trie_prefix()));
let target_chunks = targets.chunk_by_mut(move |_, next| {
if let Some(some_upper_bound) = upper_bound {
let sub_trie_prefix = next.sub_trie_prefix();
let same_chunk = sub_trie_prefix < some_upper_bound;
if !same_chunk {
upper_bound = sub_trie_upper_bound(&sub_trie_prefix);
}
same_chunk
} else {
true
}
});
// Map the chunks to the return type. Within each chunk we want targets to be sorted by their
// key, as that will be the order they are checked by the `ProofCalculator`.
target_chunks.map(move |targets| {
let prefix = targets[0].sub_trie_prefix();
let retain_root = targets[0].min_len == 0;
targets.sort_unstable_by_key(|target| target.key);
SubTrieTargets { prefix, targets, retain_root }
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_iter_sub_trie_targets() {
// Helper to create nibbles from hex string (each character is a nibble)
let nibbles = |hex: &str| -> Nibbles {
if hex.is_empty() {
return Nibbles::new();
}
format!("0x{}", hex).parse().expect("valid nibbles hex string")
};
// Test cases: (input_targets, expected_output)
// Expected output format: Vec<(exp_prefix_hex, Vec<key_hex>)>
let test_cases = vec![
// Case 1: Empty targets
(vec![], vec![]),
// Case 2: Single target without min_len
(
vec![Target::new(B256::repeat_byte(0x20))],
vec![(
"",
vec!["2020202020202020202020202020202020202020202020202020202020202020"],
)],
),
// Case 3: Multiple targets in same sub-trie (no min_len)
(
vec![Target::new(B256::repeat_byte(0x20)), Target::new(B256::repeat_byte(0x21))],
vec![(
"",
vec![
"2020202020202020202020202020202020202020202020202020202020202020",
"2121212121212121212121212121212121212121212121212121212121212121",
],
)],
),
// Case 4: Multiple targets in different sub-tries
(
vec![
Target::new(B256::repeat_byte(0x20)).with_min_len(2),
Target::new(B256::repeat_byte(0x40)).with_min_len(2),
],
vec![
("2", vec!["2020202020202020202020202020202020202020202020202020202020202020"]),
("4", vec!["4040404040404040404040404040404040404040404040404040404040404040"]),
],
),
// Case 5: Three targets, two in same sub-trie, one separate
(
vec![
Target::new(B256::repeat_byte(0x20)).with_min_len(2),
Target::new(B256::repeat_byte(0x2f)).with_min_len(2),
Target::new(B256::repeat_byte(0x40)).with_min_len(2),
],
vec![
(
"2",
vec![
"2020202020202020202020202020202020202020202020202020202020202020",
"2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f",
],
),
("4", vec!["4040404040404040404040404040404040404040404040404040404040404040"]),
],
),
// Case 6: Targets with different min_len values in same sub-trie
(
vec![
Target::new(B256::repeat_byte(0x20)).with_min_len(2),
Target::new(B256::repeat_byte(0x2f)).with_min_len(3),
],
vec![(
"2",
vec![
"2020202020202020202020202020202020202020202020202020202020202020",
"2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f",
],
)],
),
// Case 7: More complex chunking with multiple sub-tries
(
vec![
Target::new(B256::repeat_byte(0x20)).with_min_len(2),
Target::new(B256::repeat_byte(0x2f)).with_min_len(4),
Target::new(B256::repeat_byte(0x4c)).with_min_len(3),
Target::new(B256::repeat_byte(0x4c)).with_min_len(4),
Target::new(B256::repeat_byte(0x4e)).with_min_len(3),
],
vec![
(
"2",
vec![
"2020202020202020202020202020202020202020202020202020202020202020",
"2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f2f",
],
),
(
"4c",
vec![
"4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c",
"4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c4c",
],
),
(
"4e",
vec!["4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e"],
),
],
),
// Case 8: Min-len 1 should result in zero-length sub-trie prefix
(
vec![
Target::new(B256::repeat_byte(0x20)).with_min_len(1),
Target::new(B256::repeat_byte(0x40)).with_min_len(1),
],
vec![(
"",
vec![
"2020202020202020202020202020202020202020202020202020202020202020",
"4040404040404040404040404040404040404040404040404040404040404040",
],
)],
),
// Case 9: Second target's sub-trie prefix is root
(
vec![
Target::new(B256::repeat_byte(0x20)).with_min_len(2),
Target::new(B256::repeat_byte(0x40)).with_min_len(1),
],
vec![(
"",
vec![
"2020202020202020202020202020202020202020202020202020202020202020",
"4040404040404040404040404040404040404040404040404040404040404040",
],
)],
),
];
for (i, (mut input_targets, expected)) in test_cases.into_iter().enumerate() {
let test_case = i + 1;
let sub_tries: Vec<_> = iter_sub_trie_targets(&mut input_targets).collect();
assert_eq!(
sub_tries.len(),
expected.len(),
"Test case {} failed: expected {} sub-tries, got {}",
test_case,
expected.len(),
sub_tries.len()
);
for (j, (sub_trie, (exp_prefix_hex, exp_keys))) in
sub_tries.iter().zip(expected.iter()).enumerate()
{
let exp_prefix = nibbles(exp_prefix_hex);
assert_eq!(
sub_trie.prefix, exp_prefix,
"Test case {} sub-trie {}: prefix mismatch",
test_case, j
);
assert_eq!(
sub_trie.targets.len(),
exp_keys.len(),
"Test case {} sub-trie {}: expected {} targets, got {}",
test_case,
j,
exp_keys.len(),
sub_trie.targets.len()
);
for (k, (target, exp_key_hex)) in
sub_trie.targets.iter().zip(exp_keys.iter()).enumerate()
{
let exp_key = nibbles(exp_key_hex);
assert_eq!(
target.key, exp_key,
"Test case {} sub-trie {} target {}: key mismatch",
test_case, j, k
);
}
}
}
}
}

View File

@@ -123,7 +123,7 @@ where
// Compute storage root by calling storage_proof with the root path as a target.
// This returns just the root node of the storage trie.
let storage_root = storage_proof_calculator
.storage_proof(self.hashed_address, [B256::ZERO])
.storage_proof(self.hashed_address, &mut [B256::ZERO.into()])
.map(|nodes| {
// Encode the root node to RLP and hash it
let root_node =