Skip to content

Commit

Permalink
use hash builder updates
Browse files Browse the repository at this point in the history
  • Loading branch information
@shekhirin
shekhirin committed Dec 4, 2024
1 parent b074c5e commit b1424b1
Show file tree
Hide file tree
Showing 7 changed files with 138 additions and 213 deletions.
16 changes: 7 additions & 9 deletions crates/trie/parallel/src/proof.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -158,11 +158,10 @@ where

// Create a hash builder to rebuild the root node since it is not available in the database.
let retainer: ProofRetainer = targets.keys().map(Nibbles::unpack).collect();
let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);
let mut hash_builder =
HashBuilder::default().with_proof_retainer(retainer).with_updates(true);

let mut storages = HashMap::default();
let mut branch_node_hash_masks = HashMap::default();

let mut account_rlp = Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE);
let mut account_node_iter = TrieNodeIter::new(
walker,
Expand All @@ -173,7 +172,6 @@ where
{
match account_node {
TrieElement::Branch(node) => {
branch_node_hash_masks.insert(node.key.clone(), node.hash_mask);
hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
}
TrieElement::Leaf(hashed_address, account) => {
Expand Down Expand Up @@ -225,11 +223,11 @@ where
#[cfg(feature = "metrics")]
self.metrics.record_state_trie(tracker.finish());

Ok(MultiProof {
account_subtree: hash_builder.take_proof_nodes(),
storages,
branch_node_hash_masks,
})
let account_subtree = hash_builder.take_proof_nodes();
let branch_node_hash_masks =
hash_builder.split().1.into_iter().map(|(path, node)| (path, node.hash_mask)).collect();

Ok(MultiProof { account_subtree, branch_node_hash_masks, storages })
}
}

Expand Down
221 changes: 101 additions & 120 deletions crates/trie/sparse/src/state.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,7 @@ use reth_primitives_traits::Account;
use reth_tracing::tracing::trace;
use reth_trie_common::{
updates::{StorageTrieUpdates, TrieUpdates},
MultiProof, Nibbles, TrieAccount, TrieMask, TrieNode, EMPTY_ROOT_HASH,
TRIE_ACCOUNT_RLP_MAX_SIZE,
MultiProof, Nibbles, TrieAccount, TrieNode, EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE,
};
use std::{fmt, iter::Peekable};

Expand Down Expand Up @@ -103,6 +102,75 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
self.storages.get_mut(account).and_then(|e| e.as_revealed_mut())
}

/// Reveal unknown trie paths from provided leaf path and its proof for the account.
/// NOTE: This method does not extensively validate the proof.
pub fn reveal_account(
&mut self,
account: B256,
proof: impl IntoIterator<Item = (Nibbles, Bytes)>,
) -> SparseStateTrieResult<()> {
if self.is_account_revealed(&account) {
return Ok(());
}

let mut proof = proof.into_iter().peekable();

let Some(root_node) = self.validate_root_node(&mut proof)? else { return Ok(()) };

// Reveal root node if it wasn't already.
let trie = self.state.reveal_root_with_provider(
self.provider_factory.account_node_provider(),
root_node,
self.retain_updates,
)?;

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

// Mark leaf path as revealed.
self.revealed.entry(account).or_default();

Ok(())
}

/// Reveal unknown trie paths from provided leaf path and its proof for the storage slot.
/// 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<()> {
if self.is_storage_slot_revealed(&account, &slot) {
return Ok(());
}

let mut proof = proof.into_iter().peekable();

let Some(root_node) = self.validate_root_node(&mut proof)? else { return Ok(()) };

// 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,
self.retain_updates,
)?;

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

// Mark leaf path as revealed.
self.revealed.entry(account).or_default().insert(slot);

Ok(())
}

/// Reveal unknown trie paths from multiproof and the list of included accounts and slots.
/// NOTE: This method does not extensively validate the proof.
pub fn reveal_multiproof(
Expand All @@ -113,63 +181,39 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
let account_subtree = multiproof.account_subtree.into_nodes_sorted();
let mut account_nodes = account_subtree.into_iter().peekable();

if let Some((root_node, hash_mask)) =
self.validate_root_node(&mut account_nodes, &multiproof.branch_node_hash_masks)?
{
if let Some(root_node) = self.validate_root_node(&mut account_nodes)? {
// Reveal root node if it wasn't already.
let trie = self.state.reveal_root_with_provider(
self.provider_factory.account_node_provider(),
root_node,
hash_mask,
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");

let hash_mask = if let TrieNode::Branch(_) = node {
let Some(hash_mask) = multiproof.branch_node_hash_masks.get(&path) else {
return Err(SparseStateTrieError::MissingBranchNodeHashMask { path });
};
Some(*hash_mask)
} else {
None
};
trie.reveal_node(path, node, hash_mask)?;
trie.reveal_node(path, node)?;
}
}

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

if let Some((root_node, hash_mask)) =
self.validate_root_node(&mut nodes, &storage_subtree.branch_node_hash_masks)?
{
if let Some(root_node) = self.validate_root_node(&mut storage_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,
hash_mask,
self.retain_updates,
)?;

// Reveal the remaining proof nodes.
for (path, bytes) in nodes {
for (path, bytes) in storage_nodes {
let node = TrieNode::decode(&mut &bytes[..])?;
trace!(target: "trie::sparse", ?account, ?path, ?node, "Revealing storage node");
let hash_mask = if let TrieNode::Branch(_) = node {
let Some(hash_mask) = storage_subtree.branch_node_hash_masks.get(&path)
else {
return Err(SparseStateTrieError::MissingBranchNodeHashMask { path });
};
Some(*hash_mask)
} else {
None
};
trie.reveal_node(path, node, hash_mask)?;
trie.reveal_node(path, node)?;
}
}
}
Expand All @@ -185,8 +229,7 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
fn validate_root_node<I: Iterator<Item = (Nibbles, Bytes)>>(
&self,
proof: &mut Peekable<I>,
branch_node_hash_masks: &HashMap<Nibbles, TrieMask>,
) -> SparseStateTrieResult<Option<(TrieNode, Option<TrieMask>)>> {
) -> SparseStateTrieResult<Option<TrieNode>> {
let mut proof = proof.into_iter().peekable();

// Validate root node.
Expand All @@ -201,16 +244,7 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
return Err(SparseStateTrieError::InvalidRootNode { path, node })
}

let hash_mask = if matches!(root_node, TrieNode::Branch(_)) {
let Some(hash_mask) = branch_node_hash_masks.get(&path) else {
return Err(SparseStateTrieError::MissingBranchNodeHashMask { path })
};
Some(*hash_mask)
} else {
None
};

Ok(Some((root_node, hash_mask)))
Ok(Some(root_node))
}

/// Update the account leaf node.
Expand Down Expand Up @@ -360,40 +394,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, BranchNode, StorageMultiProof};
use reth_trie_common::proof::ProofRetainer;

#[test]
fn validate_root_node_first_node_not_root() {
fn validate_proof_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(), &HashMap::default()),
sparse.validate_root_node(&mut proof.into_iter().peekable()),
Err(SparseStateTrieError::InvalidRootNode { .. })
);
}

#[test]
fn validate_root_node_invalid_proof_with_empty_root() {
fn validate_proof_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(), &HashMap::default()),
Err(SparseStateTrieError::InvalidRootNode { .. })
);
}

#[test]
fn validate_root_node_no_hash_mask() {
let sparse = SparseStateTrie::default();

let mut rlp = Vec::new();
BranchNode::default().encode(&mut rlp);
let proof = [(Nibbles::default(), Bytes::from(rlp))];
assert_matches!(
sparse.validate_root_node(&mut proof.into_iter().peekable(), &HashMap::default()),
sparse.validate_root_node(&mut proof.into_iter().peekable()),
Err(SparseStateTrieError::InvalidRootNode { .. })
);
}
Expand All @@ -402,44 +423,29 @@ mod tests {
fn reveal_account_empty() {
let retainer = ProofRetainer::from_iter([Nibbles::default()]);
let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);
let root = hash_builder.root();
hash_builder.root();
let proofs = hash_builder.take_proof_nodes();
assert_eq!(proofs.len(), 1);

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

let targets = HashMap::from_iter([(root, HashSet::default())]);
let multiproof = MultiProof { account_subtree: proofs, ..Default::default() };
sparse.reveal_multiproof(targets, multiproof).unwrap();
sparse.reveal_account(Default::default(), proofs.into_inner()).unwrap();
assert_eq!(sparse.state, SparseTrie::revealed_empty());
}

#[test]
fn reveal_storage_slot_empty() {
let retainer = ProofRetainer::from_iter([Nibbles::default()]);
let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);
let root = hash_builder.root();
hash_builder.root();
let proofs = hash_builder.take_proof_nodes();
assert_eq!(proofs.len(), 1);

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

let targets =
HashMap::from_iter([(B256::default(), HashSet::from_iter([B256::default()]))]);
let multiproof = MultiProof {
storages: HashMap::from_iter([(
B256::default(),
StorageMultiProof {
root,
subtree: proofs,
branch_node_hash_masks: HashMap::default(),
},
)]),
..Default::default()
};
sparse.reveal_multiproof(targets, multiproof).unwrap();
sparse
.reveal_storage_slot(Default::default(), Default::default(), proofs.into_inner())
.unwrap();
assert_eq!(
sparse.storages,
HashMap::from_iter([(Default::default(), SparseTrie::revealed_empty())])
Expand Down Expand Up @@ -475,11 +481,9 @@ mod tests {
storage_hash_builder.add_leaf(slot_path_2.clone(), &alloy_rlp::encode_fixed_size(&value_2));

let storage_root = storage_hash_builder.root();
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 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 address_1 = b256!("1000000000000000000000000000000000000000000000000000000000000000");
let address_path_1 = Nibbles::unpack(address_1);
Expand All @@ -500,39 +504,16 @@ 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_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,
},
),
]),
},
);
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();

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

Expand Down
Loading

0 comments on commit b1424b1

Please sign in to comment.