Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(trie): use branch node hash masks in sparse trie #13135

Merged
merged 22 commits into from
Dec 6, 2024
Merged

feat(trie): use branch node hash masks in sparse trie #13135

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
70bd801
feat(trie): collect branch node hash masks when calculating a proof
shekhirin Dec 4, 2024
69d4d4e
feat(trie): collect branch node hash masks when calculating a proof
shekhirin Dec 4, 2024
b074c5e
Merge remote-tracking branch 'origin/alexey/trie-proofs-branch-node-h…
shekhirin Dec 4, 2024
b1424b1
use hash builder updates
shekhirin Dec 4, 2024
2692787
use hash builder updates
shekhirin Dec 4, 2024
d5b9b3d
Merge remote-tracking branch 'origin/alexey/trie-proofs-branch-node-h…
shekhirin Dec 4, 2024
63a6238
Revert "use hash builder updates"
shekhirin Dec 4, 2024
ee30277
fix tests
shekhirin Dec 4, 2024
139edf4
check hash mask
shekhirin Dec 4, 2024
6cf4140
do not use split
shekhirin Dec 4, 2024
1374f44
introduce a flag
shekhirin Dec 4, 2024
b25005a
extend branch node hash masks
shekhirin Dec 4, 2024
bf0d3bf
Merge remote-tracking branch 'origin/alexey/trie-proofs-branch-node-h…
shekhirin Dec 4, 2024
56314c0
add logs, remove error
shekhirin Dec 4, 2024
77c0615
revert trie node iter changes
shekhirin Dec 4, 2024
58ea42d
more logs
shekhirin Dec 4, 2024
aae8fe6
fix hash mask checking
shekhirin Dec 4, 2024
01b0ad3
Merge remote-tracking branch 'origin/main' into alexey/proofs-branch-…
shekhirin Dec 5, 2024
8fb259f
update comment
shekhirin Dec 5, 2024
88d8c56
return reveal_account and reveal_storage_slot
shekhirin Dec 5, 2024
2bb382c
add a panic comment
shekhirin Dec 5, 2024
074ebfc
add a comment about hash mask
shekhirin Dec 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
111 changes: 83 additions & 28 deletions crates/trie/sparse/src/state.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,12 +103,17 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
}

/// Reveal unknown trie paths from provided leaf path and its proof for the account.
///
/// Panics if trie updates retention is enabled.
///
/// NOTE: This method does not extensively validate the proof.
pub fn reveal_account(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i'd prefer to keep this methods, they are still useful. mb with debug_assert!(!self.updates_enabled)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fair, useful at least for tests

&mut self,
account: B256,
proof: impl IntoIterator<Item = (Nibbles, Bytes)>,
) -> SparseStateTrieResult<()> {
assert!(!self.retain_updates);

if self.is_account_revealed(&account) {
return Ok(());
}
Expand All @@ -121,13 +126,14 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
let trie = self.state.reveal_root_with_provider(
self.provider_factory.account_node_provider(),
root_node,
None,
self.retain_updates,
)?;

// Reveal the remaining proof nodes.
for (path, bytes) in proof {
let node = TrieNode::decode(&mut &bytes[..])?;
trie.reveal_node(path, node)?;
trie.reveal_node(path, node, None)?;
}

// Mark leaf path as revealed.
Expand All @@ -137,13 +143,18 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
}

/// Reveal unknown trie paths from provided leaf path and its proof for the storage slot.
///
/// Panics if trie updates retention is enabled.
///
/// NOTE: This method does not extensively validate the proof.
pub fn reveal_storage_slot(
&mut self,
account: B256,
slot: B256,
proof: impl IntoIterator<Item = (Nibbles, Bytes)>,
) -> SparseStateTrieResult<()> {
assert!(!self.retain_updates);

if self.is_storage_slot_revealed(&account, &slot) {
return Ok(());
}
Expand All @@ -156,13 +167,14 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
let trie = self.storages.entry(account).or_default().reveal_root_with_provider(
self.provider_factory.storage_node_provider(account),
root_node,
None,
self.retain_updates,
)?;

// Reveal the remaining proof nodes.
for (path, bytes) in proof {
let node = TrieNode::decode(&mut &bytes[..])?;
trie.reveal_node(path, node)?;
trie.reveal_node(path, node, None)?;
}

// Mark leaf path as revealed.
Expand All @@ -186,34 +198,48 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
let trie = self.state.reveal_root_with_provider(
self.provider_factory.account_node_provider(),
root_node,
multiproof.branch_node_hash_masks.get(&Nibbles::default()).copied(),
self.retain_updates,
)?;

// Reveal the remaining proof nodes.
for (path, bytes) in account_nodes {
let node = TrieNode::decode(&mut &bytes[..])?;
trace!(target: "trie::sparse", ?path, ?node, "Revealing account node");
trie.reveal_node(path, node)?;
let hash_mask = if let TrieNode::Branch(_) = node {
multiproof.branch_node_hash_masks.get(&path).copied()
} else {
None
};

trace!(target: "trie::sparse", ?path, ?node, ?hash_mask, "Revealing account node");
trie.reveal_node(path, node, hash_mask)?;
}
}

for (account, storage_subtree) in multiproof.storages {
let storage_subtree = storage_subtree.subtree.into_nodes_sorted();
let mut storage_nodes = storage_subtree.into_iter().peekable();
let subtree = storage_subtree.subtree.into_nodes_sorted();
let mut nodes = subtree.into_iter().peekable();

if let Some(root_node) = self.validate_root_node(&mut storage_nodes)? {
if let Some(root_node) = self.validate_root_node(&mut nodes)? {
// Reveal root node if it wasn't already.
let trie = self.storages.entry(account).or_default().reveal_root_with_provider(
self.provider_factory.storage_node_provider(account),
root_node,
storage_subtree.branch_node_hash_masks.get(&Nibbles::default()).copied(),
self.retain_updates,
)?;

// Reveal the remaining proof nodes.
for (path, bytes) in storage_nodes {
for (path, bytes) in nodes {
let node = TrieNode::decode(&mut &bytes[..])?;
trace!(target: "trie::sparse", ?account, ?path, ?node, "Revealing storage node");
trie.reveal_node(path, node)?;
let hash_mask = if let TrieNode::Branch(_) = node {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks like we could use fn is_branch()

suggestion for followup

storage_subtree.branch_node_hash_masks.get(&path).copied()
} else {
None
};

trace!(target: "trie::sparse", ?account, ?path, ?node, ?hash_mask, "Revealing storage node");
trie.reveal_node(path, node, hash_mask)?;
}
}
}
Expand Down Expand Up @@ -394,27 +420,27 @@ mod tests {
use rand::{rngs::StdRng, Rng, SeedableRng};
use reth_primitives_traits::Account;
use reth_trie::{updates::StorageTrieUpdates, HashBuilder, TrieAccount, EMPTY_ROOT_HASH};
use reth_trie_common::proof::ProofRetainer;
use reth_trie_common::{proof::ProofRetainer, StorageMultiProof, TrieMask};

#[test]
fn validate_proof_first_node_not_root() {
fn validate_root_node_first_node_not_root() {
let sparse = SparseStateTrie::default();
let proof = [(Nibbles::from_nibbles([0x1]), Bytes::from([EMPTY_STRING_CODE]))];
assert_matches!(
sparse.validate_root_node(&mut proof.into_iter().peekable()),
sparse.validate_root_node(&mut proof.into_iter().peekable(),),
Err(SparseStateTrieError::InvalidRootNode { .. })
);
}

#[test]
fn validate_proof_invalid_proof_with_empty_root() {
fn validate_root_node_invalid_proof_with_empty_root() {
let sparse = SparseStateTrie::default();
let proof = [
(Nibbles::default(), Bytes::from([EMPTY_STRING_CODE])),
(Nibbles::from_nibbles([0x1]), Bytes::new()),
];
assert_matches!(
sparse.validate_root_node(&mut proof.into_iter().peekable()),
sparse.validate_root_node(&mut proof.into_iter().peekable(),),
Err(SparseStateTrieError::InvalidRootNode { .. })
);
}
Expand All @@ -429,6 +455,7 @@ mod tests {

let mut sparse = SparseStateTrie::default();
assert_eq!(sparse.state, SparseTrie::Blind);

sparse.reveal_account(Default::default(), proofs.into_inner()).unwrap();
assert_eq!(sparse.state, SparseTrie::revealed_empty());
}
Expand All @@ -443,6 +470,7 @@ mod tests {

let mut sparse = SparseStateTrie::default();
assert!(sparse.storages.is_empty());

sparse
.reveal_storage_slot(Default::default(), Default::default(), proofs.into_inner())
.unwrap();
Expand Down Expand Up @@ -477,13 +505,15 @@ mod tests {
slot_path_1.clone(),
slot_path_2.clone(),
]));
storage_hash_builder.add_leaf(slot_path_1.clone(), &alloy_rlp::encode_fixed_size(&value_1));
storage_hash_builder.add_leaf(slot_path_2.clone(), &alloy_rlp::encode_fixed_size(&value_2));
storage_hash_builder.add_leaf(slot_path_1, &alloy_rlp::encode_fixed_size(&value_1));
storage_hash_builder.add_leaf(slot_path_2, &alloy_rlp::encode_fixed_size(&value_2));

let storage_root = storage_hash_builder.root();
let proof_nodes = storage_hash_builder.take_proof_nodes();
let storage_proof_1 = proof_nodes.matching_nodes_sorted(&slot_path_1);
let storage_proof_2 = proof_nodes.matching_nodes_sorted(&slot_path_2);
let storage_proof_nodes = storage_hash_builder.take_proof_nodes();
let storage_branch_node_hash_masks = HashMap::from_iter([
(Nibbles::default(), TrieMask::new(0b010)),
(Nibbles::from_nibbles([0x1]), TrieMask::new(0b11)),
]);

let address_1 = b256!("1000000000000000000000000000000000000000000000000000000000000000");
let address_path_1 = Nibbles::unpack(address_1);
Expand All @@ -504,16 +534,41 @@ mod tests {

let root = hash_builder.root();
let proof_nodes = hash_builder.take_proof_nodes();
let proof_1 = proof_nodes.matching_nodes_sorted(&address_path_1);
let proof_2 = proof_nodes.matching_nodes_sorted(&address_path_2);

let mut sparse = SparseStateTrie::default().with_updates(true);
sparse.reveal_account(address_1, proof_1).unwrap();
sparse.reveal_account(address_2, proof_2).unwrap();
sparse.reveal_storage_slot(address_1, slot_1, storage_proof_1.clone()).unwrap();
sparse.reveal_storage_slot(address_1, slot_2, storage_proof_2.clone()).unwrap();
sparse.reveal_storage_slot(address_2, slot_1, storage_proof_1).unwrap();
sparse.reveal_storage_slot(address_2, slot_2, storage_proof_2).unwrap();
sparse
.reveal_multiproof(
HashMap::from_iter([
(address_1, HashSet::from_iter([slot_1, slot_2])),
(address_2, HashSet::from_iter([slot_1, slot_2])),
]),
MultiProof {
account_subtree: proof_nodes,
branch_node_hash_masks: HashMap::from_iter([(
Nibbles::from_nibbles([0x1]),
TrieMask::new(0b00),
)]),
storages: HashMap::from_iter([
(
address_1,
StorageMultiProof {
root,
subtree: storage_proof_nodes.clone(),
branch_node_hash_masks: storage_branch_node_hash_masks.clone(),
},
),
(
address_2,
StorageMultiProof {
root,
subtree: storage_proof_nodes,
branch_node_hash_masks: storage_branch_node_hash_masks,
},
),
]),
},
)
.unwrap();

assert_eq!(sparse.root(), Some(root));

Expand Down
Loading
Loading