use crate::{walker::TrieWalker, BranchNodeCompact, HashBuilder, Nibbles};
use alloy_primitives::B256;
use std::collections::{HashMap, HashSet};
#[derive(PartialEq, Eq, Clone, Default, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct TrieUpdates {
pub(crate) account_nodes: HashMap<Nibbles, BranchNodeCompact>,
pub(crate) removed_nodes: HashSet<Nibbles>,
pub(crate) storage_tries: HashMap<B256, StorageTrieUpdates>,
}
impl TrieUpdates {
pub fn is_empty(&self) -> bool {
self.account_nodes.is_empty() &&
self.removed_nodes.is_empty() &&
self.storage_tries.is_empty()
}
pub const fn account_nodes_ref(&self) -> &HashMap<Nibbles, BranchNodeCompact> {
&self.account_nodes
}
pub const fn removed_nodes_ref(&self) -> &HashSet<Nibbles> {
&self.removed_nodes
}
pub const fn storage_tries_ref(&self) -> &HashMap<B256, StorageTrieUpdates> {
&self.storage_tries
}
pub fn extend(&mut self, other: Self) {
self.extend_common(&other);
self.account_nodes.extend(exclude_empty_from_pair(other.account_nodes));
self.removed_nodes.extend(exclude_empty(other.removed_nodes));
for (hashed_address, storage_trie) in other.storage_tries {
self.storage_tries.entry(hashed_address).or_default().extend(storage_trie);
}
}
pub fn extend_ref(&mut self, other: &Self) {
self.extend_common(other);
self.account_nodes.extend(exclude_empty_from_pair(
other.account_nodes.iter().map(|(k, v)| (k.clone(), v.clone())),
));
self.removed_nodes.extend(exclude_empty(other.removed_nodes.iter().cloned()));
for (hashed_address, storage_trie) in &other.storage_tries {
self.storage_tries.entry(*hashed_address).or_default().extend_ref(storage_trie);
}
}
fn extend_common(&mut self, other: &Self) {
self.account_nodes.retain(|nibbles, _| !other.removed_nodes.contains(nibbles));
}
pub fn insert_storage_updates(
&mut self,
hashed_address: B256,
storage_updates: StorageTrieUpdates,
) {
let existing = self.storage_tries.insert(hashed_address, storage_updates);
debug_assert!(existing.is_none());
}
pub fn finalize<C>(
&mut self,
walker: TrieWalker<C>,
hash_builder: HashBuilder,
destroyed_accounts: HashSet<B256>,
) {
let (_, removed_node_keys) = walker.split();
self.removed_nodes.extend(exclude_empty(removed_node_keys));
let (_, updated_nodes) = hash_builder.split();
self.account_nodes.extend(exclude_empty_from_pair(updated_nodes));
for destroyed in destroyed_accounts {
self.storage_tries.entry(destroyed).or_default().set_deleted(true);
}
}
pub fn into_sorted(self) -> TrieUpdatesSorted {
let mut account_nodes = Vec::from_iter(self.account_nodes);
account_nodes.sort_unstable_by(|a, b| a.0.cmp(&b.0));
let storage_tries = self
.storage_tries
.into_iter()
.map(|(hashed_address, updates)| (hashed_address, updates.into_sorted()))
.collect();
TrieUpdatesSorted { removed_nodes: self.removed_nodes, account_nodes, storage_tries }
}
}
#[derive(PartialEq, Eq, Clone, Default, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct StorageTrieUpdates {
pub(crate) is_deleted: bool,
pub(crate) storage_nodes: HashMap<Nibbles, BranchNodeCompact>,
pub(crate) removed_nodes: HashSet<Nibbles>,
}
#[cfg(feature = "test-utils")]
impl StorageTrieUpdates {
pub fn new(updates: impl IntoIterator<Item = (Nibbles, BranchNodeCompact)>) -> Self {
Self { storage_nodes: exclude_empty_from_pair(updates).collect(), ..Default::default() }
}
}
impl StorageTrieUpdates {
pub fn deleted() -> Self {
Self {
is_deleted: true,
storage_nodes: HashMap::default(),
removed_nodes: HashSet::default(),
}
}
pub fn len(&self) -> usize {
(self.is_deleted as usize) + self.storage_nodes.len() + self.removed_nodes.len()
}
pub const fn is_deleted(&self) -> bool {
self.is_deleted
}
pub const fn storage_nodes_ref(&self) -> &HashMap<Nibbles, BranchNodeCompact> {
&self.storage_nodes
}
pub const fn removed_nodes_ref(&self) -> &HashSet<Nibbles> {
&self.removed_nodes
}
pub fn is_empty(&self) -> bool {
!self.is_deleted && self.storage_nodes.is_empty() && self.removed_nodes.is_empty()
}
pub fn set_deleted(&mut self, deleted: bool) {
self.is_deleted = deleted;
}
pub fn extend(&mut self, other: Self) {
self.extend_common(&other);
self.storage_nodes.extend(exclude_empty_from_pair(other.storage_nodes));
self.removed_nodes.extend(exclude_empty(other.removed_nodes));
}
pub fn extend_ref(&mut self, other: &Self) {
self.extend_common(other);
self.storage_nodes.extend(exclude_empty_from_pair(
other.storage_nodes.iter().map(|(k, v)| (k.clone(), v.clone())),
));
self.removed_nodes.extend(exclude_empty(other.removed_nodes.iter().cloned()));
}
fn extend_common(&mut self, other: &Self) {
if other.is_deleted {
self.storage_nodes.clear();
self.removed_nodes.clear();
}
self.is_deleted |= other.is_deleted;
self.storage_nodes.retain(|nibbles, _| !other.removed_nodes.contains(nibbles));
}
pub fn finalize<C>(&mut self, walker: TrieWalker<C>, hash_builder: HashBuilder) {
let (_, removed_keys) = walker.split();
self.removed_nodes.extend(exclude_empty(removed_keys));
let (_, updated_nodes) = hash_builder.split();
self.storage_nodes.extend(exclude_empty_from_pair(updated_nodes));
}
pub fn into_sorted(self) -> StorageTrieUpdatesSorted {
let mut storage_nodes = Vec::from_iter(self.storage_nodes);
storage_nodes.sort_unstable_by(|a, b| a.0.cmp(&b.0));
StorageTrieUpdatesSorted {
is_deleted: self.is_deleted,
removed_nodes: self.removed_nodes,
storage_nodes,
}
}
}
#[derive(PartialEq, Eq, Clone, Default, Debug)]
pub struct TrieUpdatesSorted {
pub(crate) account_nodes: Vec<(Nibbles, BranchNodeCompact)>,
pub(crate) removed_nodes: HashSet<Nibbles>,
pub(crate) storage_tries: HashMap<B256, StorageTrieUpdatesSorted>,
}
impl TrieUpdatesSorted {
pub fn account_nodes_ref(&self) -> &[(Nibbles, BranchNodeCompact)] {
&self.account_nodes
}
pub const fn removed_nodes_ref(&self) -> &HashSet<Nibbles> {
&self.removed_nodes
}
pub const fn storage_tries_ref(&self) -> &HashMap<B256, StorageTrieUpdatesSorted> {
&self.storage_tries
}
}
#[derive(PartialEq, Eq, Clone, Default, Debug)]
pub struct StorageTrieUpdatesSorted {
pub(crate) is_deleted: bool,
pub(crate) storage_nodes: Vec<(Nibbles, BranchNodeCompact)>,
pub(crate) removed_nodes: HashSet<Nibbles>,
}
impl StorageTrieUpdatesSorted {
pub const fn is_deleted(&self) -> bool {
self.is_deleted
}
pub fn storage_nodes_ref(&self) -> &[(Nibbles, BranchNodeCompact)] {
&self.storage_nodes
}
pub const fn removed_nodes_ref(&self) -> &HashSet<Nibbles> {
&self.removed_nodes
}
}
fn exclude_empty(iter: impl IntoIterator<Item = Nibbles>) -> impl Iterator<Item = Nibbles> {
iter.into_iter().filter(|n| !n.is_empty())
}
fn exclude_empty_from_pair<V>(
iter: impl IntoIterator<Item = (Nibbles, V)>,
) -> impl Iterator<Item = (Nibbles, V)> {
iter.into_iter().filter(|(n, _)| !n.is_empty())
}