jleony/nanoreth
1
0
Fork 0
mirror of https://github.com/hl-archive-node/nanoreth.git synced 2025-12-06 10:59:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(trie): sparse trie methods for trie task integration (#12720)

Browse Source
This commit is contained in:
Alexey Shekhirin
2024-11-25 13:13:01 +00:00
committed by GitHub
parent caac226c73
commit 04dd005af9
5 changed files with 344 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
Cargo.lock generated
View File

@ -9454,6 +9454,7 @@ version = "1.1.2"
dependencies = [
"alloy-primitives",
"alloy-rlp",
"arbitrary",
"assert_matches",
"criterion",
"itertools 0.13.0",

1
crates/trie/sparse/Cargo.toml
View File

@ -32,6 +32,7 @@ reth-testing-utils.workspace = true
reth-trie = { workspace = true, features = ["test-utils"] }
reth-trie-common = { workspace = true, features = ["test-utils", "arbitrary"] }
arbitrary.workspace = true
assert_matches.workspace = true
criterion.workspace = true
itertools.workspace = true

263
crates/trie/sparse/src/state.rs
View File

@ -6,17 +6,23 @@ use alloy_primitives::{
Bytes, B256,
};
use alloy_rlp::Decodable;
use reth_trie::{Nibbles, TrieNode};
use reth_trie::{
updates::{StorageTrieUpdates, TrieUpdates},
Nibbles, TrieNode,
};
/// Sparse state trie representing lazy-loaded Ethereum state trie.
#[derive(Default, Debug)]
pub struct SparseStateTrie {
retain_updates: bool,
/// Sparse account trie.
pub(crate) state: SparseTrie,
state: SparseTrie,
/// Sparse storage tries.
pub(crate) storages: HashMap<B256, SparseTrie>,
storages: HashMap<B256, SparseTrie>,
/// Collection of revealed account and storage keys.
pub(crate) revealed: HashMap<B256, HashSet<B256>>,
revealed: HashMap<B256, HashSet<B256>>,
/// Collection of addresses that had their storage tries wiped.
wiped_storages: HashSet<B256>,
}
impl SparseStateTrie {
@ -25,6 +31,12 @@ impl SparseStateTrie {
Self { state, ..Default::default() }
}
/// Set the retention of branch node updates and deletions.
pub const fn with_updates(mut self, retain_updates: bool) -> Self {
self.retain_updates = retain_updates;
self
}
/// Returns `true` if account was already revealed.
pub fn is_account_revealed(&self, account: &B256) -> bool {
self.revealed.contains_key(account)
@ -42,7 +54,7 @@ impl SparseStateTrie {
account: B256,
proof: impl IntoIterator<Item = (Nibbles, Bytes)>,
) -> SparseStateTrieResult<()> {
if self.revealed.contains_key(&account) {
if self.is_account_revealed(&account) {
return Ok(());
}
@ -51,7 +63,7 @@ impl SparseStateTrie {
let Some(root_node) = self.validate_proof(&mut proof)? else { return Ok(()) };
// Reveal root node if it wasn't already.
let trie = self.state.reveal_root(root_node)?;
let trie = self.state.reveal_root(root_node, self.retain_updates)?;
// Reveal the remaining proof nodes.
for (path, bytes) in proof {
@ -73,7 +85,7 @@ impl SparseStateTrie {
slot: B256,
proof: impl IntoIterator<Item = (Nibbles, Bytes)>,
) -> SparseStateTrieResult<()> {
if self.revealed.get(&account).is_some_and(|v| v.contains(&slot)) {
if self.is_storage_slot_revealed(&account, &slot) {
return Ok(());
}
@ -82,7 +94,11 @@ impl SparseStateTrie {
let Some(root_node) = self.validate_proof(&mut proof)? else { return Ok(()) };
// Reveal root node if it wasn't already.
let trie = self.storages.entry(account).or_default().reveal_root(root_node)?;
let trie = self
.storages
.entry(account)
.or_default()
.reveal_root(root_node, self.retain_updates)?;
// Reveal the remaining proof nodes.
for (path, bytes) in proof {
@ -118,30 +134,98 @@ impl SparseStateTrie {
Ok(Some(root_node))
}
/// Update the leaf node.
pub fn update_leaf(&mut self, path: Nibbles, value: Vec<u8>) -> SparseStateTrieResult<()> {
/// Update the account leaf node.
pub fn update_account_leaf(
&mut self,
path: Nibbles,
value: Vec<u8>,
) -> SparseStateTrieResult<()> {
self.state.update_leaf(path, value)?;
Ok(())
}
/// Remove the account leaf node.
pub fn remove_account_leaf(&mut self, path: &Nibbles) -> SparseStateTrieResult<()> {
self.state.remove_leaf(path)?;
Ok(())
}
/// Returns sparse trie root if the trie has been revealed.
pub fn root(&mut self) -> Option<B256> {
self.state.root()
}
/// Calculates the hashes of the nodes below the provided level.
pub fn calculate_below_level(&mut self, level: usize) {
self.state.calculate_below_level(level);
}
/// Update the leaf node of a storage trie at the provided address.
pub fn update_storage_leaf(
&mut self,
address: B256,
slot: Nibbles,
value: Vec<u8>,
) -> SparseStateTrieResult<()> {
self.storages.entry(address).or_default().update_leaf(slot, value)?;
Ok(())
}
/// Wipe the storage trie at the provided address.
pub fn wipe_storage(&mut self, address: B256) -> SparseStateTrieResult<()> {
let Some(trie) = self.storages.get_mut(&address) else { return Ok(()) };
self.wiped_storages.insert(address);
trie.wipe().map_err(Into::into)
}
/// Returns storage sparse trie root if the trie has been revealed.
pub fn storage_root(&mut self, account: B256) -> Option<B256> {
self.storages.get_mut(&account).and_then(|trie| trie.root())
}
/// Returns [`TrieUpdates`] by taking the updates from the revealed sparse tries.
///
/// Returns `None` if the accounts trie is not revealed.
pub fn take_trie_updates(&mut self) -> Option<TrieUpdates> {
self.state.as_revealed_mut().map(|state| {
let updates = state.take_updates();
TrieUpdates {
account_nodes: HashMap::from_iter(updates.updated_nodes),
removed_nodes: HashSet::from_iter(updates.removed_nodes),
storage_tries: self
.storages
.iter_mut()
.map(|(address, trie)| {
let trie = trie.as_revealed_mut().unwrap();
let updates = trie.take_updates();
let updates = StorageTrieUpdates {
is_deleted: self.wiped_storages.contains(address),
storage_nodes: HashMap::from_iter(updates.updated_nodes),
removed_nodes: HashSet::from_iter(updates.removed_nodes),
};
(*address, updates)
})
.filter(|(_, updates)| !updates.is_empty())
.collect(),
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use alloy_primitives::Bytes;
use alloy_primitives::{b256, Bytes, U256};
use alloy_rlp::EMPTY_STRING_CODE;
use arbitrary::Arbitrary;
use assert_matches::assert_matches;
use reth_trie::HashBuilder;
use itertools::Itertools;
use rand::{rngs::StdRng, Rng, SeedableRng};
use reth_primitives_traits::Account;
use reth_trie::{
updates::StorageTrieUpdates, BranchNodeCompact, HashBuilder, TrieAccount, TrieMask,
EMPTY_ROOT_HASH,
};
use reth_trie_common::proof::ProofRetainer;
#[test]
@ -199,4 +283,159 @@ mod tests {
HashMap::from_iter([(Default::default(), SparseTrie::revealed_empty())])
);
}
#[test]
fn take_trie_updates() {
reth_tracing::init_test_tracing();
// let mut rng = generators::rng();
let mut rng = StdRng::seed_from_u64(1);
let mut bytes = [0u8; 1024];
rng.fill(bytes.as_mut_slice());
let slot_1 = b256!("1000000000000000000000000000000000000000000000000000000000000000");
let slot_path_1 = Nibbles::unpack(slot_1);
let value_1 = U256::from(rng.gen::<u64>());
let slot_2 = b256!("1100000000000000000000000000000000000000000000000000000000000000");
let slot_path_2 = Nibbles::unpack(slot_2);
let value_2 = U256::from(rng.gen::<u64>());
let slot_3 = b256!("2000000000000000000000000000000000000000000000000000000000000000");
let slot_path_3 = Nibbles::unpack(slot_3);
let value_3 = U256::from(rng.gen::<u64>());
let mut storage_hash_builder =
HashBuilder::default().with_proof_retainer(ProofRetainer::from_iter([
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));
let storage_root = storage_hash_builder.root();
let proof_nodes = storage_hash_builder.take_proof_nodes();
let storage_proof_1 = proof_nodes
.iter()
.filter(|(path, _)| path.is_empty() || slot_path_1.common_prefix_length(path) > 0)
.map(|(path, proof)| (path.clone(), proof.clone()))
.sorted_by_key(|(path, _)| path.clone())
.collect::<Vec<_>>();
let storage_proof_2 = proof_nodes
.iter()
.filter(|(path, _)| path.is_empty() || slot_path_2.common_prefix_length(path) > 0)
.map(|(path, proof)| (path.clone(), proof.clone()))
.sorted_by_key(|(path, _)| path.clone())
.collect::<Vec<_>>();
let address_1 = b256!("1000000000000000000000000000000000000000000000000000000000000000");
let address_path_1 = Nibbles::unpack(address_1);
let account_1 = Account::arbitrary(&mut arbitrary::Unstructured::new(&bytes)).unwrap();
let mut trie_account_1 = TrieAccount::from((account_1, storage_root));
let address_2 = b256!("1100000000000000000000000000000000000000000000000000000000000000");
let address_path_2 = Nibbles::unpack(address_2);
let account_2 = Account::arbitrary(&mut arbitrary::Unstructured::new(&bytes)).unwrap();
let mut trie_account_2 = TrieAccount::from((account_2, EMPTY_ROOT_HASH));
let mut hash_builder =
HashBuilder::default().with_proof_retainer(ProofRetainer::from_iter([
address_path_1.clone(),
address_path_2.clone(),
]));
hash_builder.add_leaf(address_path_1.clone(), &alloy_rlp::encode(trie_account_1));
hash_builder.add_leaf(address_path_2.clone(), &alloy_rlp::encode(trie_account_2));
let root = hash_builder.root();
let proof_nodes = hash_builder.take_proof_nodes();
let proof_1 = proof_nodes
.iter()
.filter(|(path, _)| path.is_empty() || address_path_1.common_prefix_length(path) > 0)
.map(|(path, proof)| (path.clone(), proof.clone()))
.sorted_by_key(|(path, _)| path.clone())
.collect::<Vec<_>>();
let proof_2 = proof_nodes
.iter()
.filter(|(path, _)| path.is_empty() || address_path_2.common_prefix_length(path) > 0)
.map(|(path, proof)| (path.clone(), proof.clone()))
.sorted_by_key(|(path, _)| path.clone())
.collect::<Vec<_>>();
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();
assert_eq!(sparse.root(), Some(root));
let address_3 = b256!("2000000000000000000000000000000000000000000000000000000000000000");
let address_path_3 = Nibbles::unpack(address_3);
let account_3 = Account { nonce: account_1.nonce + 1, ..account_1 };
let trie_account_3 = TrieAccount::from((account_3, EMPTY_ROOT_HASH));
sparse.update_account_leaf(address_path_3, alloy_rlp::encode(trie_account_3)).unwrap();
sparse.update_storage_leaf(address_1, slot_path_3, alloy_rlp::encode(value_3)).unwrap();
trie_account_1.storage_root = sparse.storage_root(address_1).unwrap();
sparse.update_account_leaf(address_path_1, alloy_rlp::encode(trie_account_1)).unwrap();
sparse.wipe_storage(address_2).unwrap();
trie_account_2.storage_root = sparse.storage_root(address_2).unwrap();
sparse.update_account_leaf(address_path_2, alloy_rlp::encode(trie_account_2)).unwrap();
sparse.root();
let sparse_updates = sparse.take_trie_updates().unwrap();
// TODO(alexey): assert against real state root calculation updates
pretty_assertions::assert_eq!(
sparse_updates,
TrieUpdates {
account_nodes: HashMap::from_iter([
(
Nibbles::default(),
BranchNodeCompact {
state_mask: TrieMask::new(0b110),
tree_mask: TrieMask::new(0b000),
hash_mask: TrieMask::new(0b010),
hashes: vec![b256!(
"4c4ffbda3569fcf2c24ea2000b4cec86ef8b92cbf9ff415db43184c0f75a212e"
)],
root_hash: Some(b256!(
"60944bd29458529c3065d19f63c6e3d5269596fd3b04ca2e7b318912dc89ca4c"
))
},
),
]),
storage_tries: HashMap::from_iter([
(
b256!("1000000000000000000000000000000000000000000000000000000000000000"),
StorageTrieUpdates {
is_deleted: false,
storage_nodes: HashMap::from_iter([(
Nibbles::default(),
BranchNodeCompact {
state_mask: TrieMask::new(0b110),
tree_mask: TrieMask::new(0b000),
hash_mask: TrieMask::new(0b010),
hashes: vec![b256!("5bc8b4fdf51839c1e18b8d6a4bd3e2e52c9f641860f0e4d197b68c2679b0e436")],
root_hash: Some(b256!("c44abf1a9e1a92736ac479b20328e8d7998aa8838b6ef52620324c9ce85e3201"))
}
)]),
removed_nodes: HashSet::default()
}
),
(
b256!("1100000000000000000000000000000000000000000000000000000000000000"),
StorageTrieUpdates {
is_deleted: true,
storage_nodes: HashMap::default(),
removed_nodes: HashSet::default()
}
)
]),
removed_nodes: HashSet::default()
}
);
}
}

92
crates/trie/sparse/src/trie.rs
View File

@ -53,9 +53,13 @@ impl SparseTrie {
/// # Returns
///
/// Mutable reference to [`RevealedSparseTrie`].
pub fn reveal_root(&mut self, root: TrieNode) -> SparseTrieResult<&mut RevealedSparseTrie> {
pub fn reveal_root(
&mut self,
root: TrieNode,
retain_updates: bool,
) -> SparseTrieResult<&mut RevealedSparseTrie> {
if self.is_blind() {
*self = Self::Revealed(Box::new(RevealedSparseTrie::from_root(root)?))
*self = Self::Revealed(Box::new(RevealedSparseTrie::from_root(root, retain_updates)?))
}
Ok(self.as_revealed_mut().unwrap())
}
@ -67,10 +71,29 @@ impl SparseTrie {
Ok(())
}
/// Remove the leaf node.
pub fn remove_leaf(&mut self, path: &Nibbles) -> SparseTrieResult<()> {
let revealed = self.as_revealed_mut().ok_or(SparseTrieError::Blind)?;
revealed.remove_leaf(path)?;
Ok(())
}
/// Wipe the trie, removing all values and nodes, and replacing the root with an empty node.
pub fn wipe(&mut self) -> SparseTrieResult<()> {
let revealed = self.as_revealed_mut().ok_or(SparseTrieError::Blind)?;
revealed.wipe();
Ok(())
}
/// Calculates and returns the trie root if the trie has been revealed.
pub fn root(&mut self) -> Option<B256> {
Some(self.as_revealed_mut()?.root())
}
/// Calculates the hashes of the nodes below the provided level.
pub fn calculate_below_level(&mut self, level: usize) {
self.as_revealed_mut().unwrap().update_rlp_node_level(level);
}
}
/// The representation of revealed sparse trie.
@ -120,19 +143,20 @@ impl Default for RevealedSparseTrie {
impl RevealedSparseTrie {
/// Create new revealed sparse trie from the given root node.
pub fn from_root(node: TrieNode) -> SparseTrieResult<Self> {
pub fn from_root(node: TrieNode, retain_updates: bool) -> SparseTrieResult<Self> {
let mut this = Self {
nodes: HashMap::default(),
values: HashMap::default(),
prefix_set: PrefixSetMut::default(),
rlp_buf: Vec::new(),
updates: None,
};
}
.with_updates(retain_updates);
this.reveal_node(Nibbles::default(), node)?;
Ok(this)
}
/// Makes the sparse trie to store updated branch nodes.
/// Set the retention of branch node updates and deletions.
pub fn with_updates(mut self, retain_updates: bool) -> Self {
if retain_updates {
self.updates = Some(SparseTrieUpdates::default());
@ -580,6 +604,12 @@ impl RevealedSparseTrie {
Ok(nodes)
}
/// Wipe the trie, removing all values and nodes, and replacing the root with an empty node.
pub fn wipe(&mut self) {
*self = Self::default();
self.prefix_set = PrefixSetMut::all();
}
/// Return the root of the sparse trie.
/// Updates all remaining dirty nodes before calculating the root.
pub fn root(&mut self) -> B256 {
@ -773,8 +803,7 @@ impl RevealedSparseTrie {
}
// Set the hash mask. If a child node has a hash value AND is a
// branch node, set the hash mask
// and save the hash.
// branch node, set the hash mask and save the hash.
let hash = child.as_hash().filter(|_| node_type.is_branch());
hash_mask_values.push(hash.is_some());
if let Some(hash) = hash {
@ -998,8 +1027,8 @@ impl RlpNodeBuffers {
/// The aggregation of sparse trie updates.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct SparseTrieUpdates {
updated_nodes: HashMap<Nibbles, BranchNodeCompact>,
removed_nodes: HashSet<Nibbles>,
pub(crate) updated_nodes: HashMap<Nibbles, BranchNodeCompact>,
pub(crate) removed_nodes: HashSet<Nibbles>,
}
#[cfg(test)]
@ -1560,7 +1589,7 @@ mod tests {
TrieMask::new(0b11),
));
let mut sparse = RevealedSparseTrie::from_root(branch.clone()).unwrap();
let mut sparse = RevealedSparseTrie::from_root(branch.clone(), false).unwrap();
// Reveal a branch node and one of its children
//
@ -1722,6 +1751,7 @@ mod tests {
.take_proof_nodes();
let mut sparse = RevealedSparseTrie::from_root(
TrieNode::decode(&mut &proof_nodes.nodes_sorted()[0].1[..]).unwrap(),
false,
)
.unwrap();
@ -1796,6 +1826,7 @@ mod tests {
.take_proof_nodes();
let mut sparse = RevealedSparseTrie::from_root(
TrieNode::decode(&mut &proof_nodes.nodes_sorted()[0].1[..]).unwrap(),
false,
)
.unwrap();
@ -1866,6 +1897,7 @@ mod tests {
.take_proof_nodes();
let mut sparse = RevealedSparseTrie::from_root(
TrieNode::decode(&mut &proof_nodes.nodes_sorted()[0].1[..]).unwrap(),
false,
)
.unwrap();
@ -1993,4 +2025,44 @@ mod tests {
assert_eq!(sparse_root, hash_builder.root());
assert_eq!(sparse_updates.updated_nodes, hash_builder.updated_branch_nodes.take().unwrap());
}
#[test]
fn sparse_trie_wipe() {
let mut sparse = RevealedSparseTrie::default().with_updates(true);
let value = alloy_rlp::encode_fixed_size(&U256::ZERO).to_vec();
// Extension (Key = 5) – Level 0
// └── Branch (Mask = 1011) – Level 1
// ├── 0 -> Extension (Key = 23) – Level 2
// │ └── Branch (Mask = 0101) – Level 3
// │ ├── 1 -> Leaf (Key = 1, Path = 50231) – Level 4
// │ └── 3 -> Leaf (Key = 3, Path = 50233) – Level 4
// ├── 2 -> Leaf (Key = 013, Path = 52013) – Level 2
// └── 3 -> Branch (Mask = 0101) – Level 2
// ├── 1 -> Leaf (Key = 3102, Path = 53102) – Level 3
// └── 3 -> Branch (Mask = 1010) – Level 3
// ├── 0 -> Leaf (Key = 3302, Path = 53302) – Level 4
// └── 2 -> Leaf (Key = 3320, Path = 53320) – Level 4
sparse
.update_leaf(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]), value.clone())
.unwrap();
sparse
.update_leaf(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]), value.clone())
.unwrap();
sparse
.update_leaf(Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3]), value.clone())
.unwrap();
sparse
.update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2]), value.clone())
.unwrap();
sparse
.update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2]), value.clone())
.unwrap();
sparse.update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0]), value).unwrap();
sparse.wipe();
assert_eq!(sparse.root(), EMPTY_ROOT_HASH);
}
}

15
crates/trie/trie/src/updates.rs
View File

@ -6,11 +6,14 @@ use std::collections::{HashMap, HashSet};
#[derive(PartialEq, Eq, Clone, Default, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct TrieUpdates {
/// Collection of updated intermediate account nodes indexed by full path.
#[cfg_attr(feature = "serde", serde(with = "serde_nibbles_map"))]
pub(crate) account_nodes: HashMap<Nibbles, BranchNodeCompact>,
pub account_nodes: HashMap<Nibbles, BranchNodeCompact>,
/// Collection of removed intermediate account nodes indexed by full path.
#[cfg_attr(feature = "serde", serde(with = "serde_nibbles_set"))]
pub(crate) removed_nodes: HashSet<Nibbles>,
pub(crate) storage_tries: HashMap<B256, StorageTrieUpdates>,
pub removed_nodes: HashSet<Nibbles>,
/// Collection of updated storage tries indexed by the hashed address.
pub storage_tries: HashMap<B256, StorageTrieUpdates>,
}
impl TrieUpdates {
@ -113,13 +116,13 @@ impl TrieUpdates {
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct StorageTrieUpdates {
/// Flag indicating whether the trie was deleted.
pub(crate) is_deleted: bool,
pub is_deleted: bool,
/// Collection of updated storage trie nodes.
#[cfg_attr(feature = "serde", serde(with = "serde_nibbles_map"))]
pub(crate) storage_nodes: HashMap<Nibbles, BranchNodeCompact>,
pub storage_nodes: HashMap<Nibbles, BranchNodeCompact>,
/// Collection of removed storage trie nodes.
#[cfg_attr(feature = "serde", serde(with = "serde_nibbles_set"))]
pub(crate) removed_nodes: HashSet<Nibbles>,
pub removed_nodes: HashSet<Nibbles>,
}
#[cfg(feature = "test-utils")]