Skip to main content

reth_trie_common/
updates.rs

1use crate::{
2    utils::{extend_sorted_vec, kway_merge_sorted},
3    BranchNodeCompact, HashBuilder, Nibbles,
4};
5use alloc::{
6    collections::{btree_map::BTreeMap, btree_set::BTreeSet},
7    vec::Vec,
8};
9use alloy_primitives::{
10    map::{B256Map, B256Set, HashMap, HashSet},
11    FixedBytes, B256,
12};
13
14/// The aggregation of trie updates.
15#[derive(PartialEq, Eq, Clone, Default, Debug)]
16#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
17pub struct TrieUpdates {
18    /// Collection of updated intermediate account nodes indexed by full path.
19    #[cfg_attr(any(test, feature = "serde"), serde(with = "serde_nibbles_map"))]
20    pub account_nodes: HashMap<Nibbles, BranchNodeCompact>,
21    /// Collection of removed intermediate account nodes indexed by full path.
22    #[cfg_attr(any(test, feature = "serde"), serde(with = "serde_nibbles_set"))]
23    pub removed_nodes: HashSet<Nibbles>,
24    /// Collection of updated storage tries indexed by the hashed address.
25    pub storage_tries: B256Map<StorageTrieUpdates>,
26}
27
28impl TrieUpdates {
29    /// Creates a new `TrieUpdates` with pre-allocated capacity.
30    pub fn with_capacity(account_nodes: usize, storage_tries: usize) -> Self {
31        Self {
32            account_nodes: HashMap::with_capacity_and_hasher(account_nodes, Default::default()),
33            removed_nodes: HashSet::with_capacity_and_hasher(account_nodes / 4, Default::default()),
34            storage_tries: B256Map::with_capacity_and_hasher(storage_tries, Default::default()),
35        }
36    }
37
38    /// Returns `true` if the updates are empty.
39    pub fn is_empty(&self) -> bool {
40        self.account_nodes.is_empty() &&
41            self.removed_nodes.is_empty() &&
42            self.storage_tries.is_empty()
43    }
44
45    /// Returns reference to updated account nodes.
46    pub const fn account_nodes_ref(&self) -> &HashMap<Nibbles, BranchNodeCompact> {
47        &self.account_nodes
48    }
49
50    /// Returns a reference to removed account nodes.
51    pub const fn removed_nodes_ref(&self) -> &HashSet<Nibbles> {
52        &self.removed_nodes
53    }
54
55    /// Returns a reference to updated storage tries.
56    pub const fn storage_tries_ref(&self) -> &B256Map<StorageTrieUpdates> {
57        &self.storage_tries
58    }
59
60    /// Extends the trie updates.
61    pub fn extend(&mut self, other: Self) {
62        self.extend_common(&other);
63        self.account_nodes.extend(exclude_empty_from_pair(other.account_nodes));
64        self.removed_nodes.extend(exclude_empty(other.removed_nodes));
65        for (hashed_address, storage_trie) in other.storage_tries {
66            self.storage_tries.entry(hashed_address).or_default().extend(storage_trie);
67        }
68    }
69
70    /// Extends the trie updates.
71    ///
72    /// Slightly less efficient than [`Self::extend`], but preferred to `extend(other.clone())`.
73    pub fn extend_ref(&mut self, other: &Self) {
74        self.extend_common(other);
75        self.account_nodes.extend(exclude_empty_from_pair(
76            other.account_nodes.iter().map(|(k, v)| (*k, v.clone())),
77        ));
78        self.removed_nodes.extend(exclude_empty(other.removed_nodes.iter().copied()));
79        for (hashed_address, storage_trie) in &other.storage_tries {
80            self.storage_tries.entry(*hashed_address).or_default().extend_ref(storage_trie);
81        }
82    }
83
84    fn extend_common(&mut self, other: &Self) {
85        self.account_nodes.retain(|nibbles, _| !other.removed_nodes.contains(nibbles));
86    }
87
88    /// Extend trie updates with sorted data, converting directly into the unsorted `HashMap`
89    /// representation. This is more efficient than first converting to `TrieUpdates` and
90    /// then extending, as it avoids creating intermediate `HashMap` allocations.
91    ///
92    /// This top-level helper merges account nodes and delegates each account's storage trie to
93    /// [`StorageTrieUpdates::extend_from_sorted`].
94    pub fn extend_from_sorted(&mut self, sorted: &TrieUpdatesSorted) {
95        // Reserve capacity for account nodes
96        let new_nodes_count = sorted.account_nodes.len();
97        self.account_nodes.reserve(new_nodes_count);
98
99        // Insert account nodes from sorted (only non-None entries)
100        for (nibbles, maybe_node) in &sorted.account_nodes {
101            if nibbles.is_empty() {
102                continue;
103            }
104            match maybe_node {
105                Some(node) => {
106                    self.removed_nodes.remove(nibbles);
107                    self.account_nodes.insert(*nibbles, node.clone());
108                }
109                None => {
110                    self.account_nodes.remove(nibbles);
111                    self.removed_nodes.insert(*nibbles);
112                }
113            }
114        }
115
116        // Extend storage tries
117        self.storage_tries.reserve(sorted.storage_tries.len());
118        for (hashed_address, sorted_storage) in &sorted.storage_tries {
119            self.storage_tries
120                .entry(*hashed_address)
121                .or_default()
122                .extend_from_sorted(sorted_storage);
123        }
124    }
125
126    /// Insert storage updates for a given hashed address.
127    pub fn insert_storage_updates(
128        &mut self,
129        hashed_address: B256,
130        storage_updates: StorageTrieUpdates,
131    ) {
132        if storage_updates.is_empty() {
133            return;
134        }
135        let existing = self.storage_tries.insert(hashed_address, storage_updates);
136        debug_assert!(existing.is_none());
137    }
138
139    /// Finalize state trie updates.
140    pub fn finalize(
141        &mut self,
142        hash_builder: HashBuilder,
143        removed_keys: HashSet<Nibbles>,
144        destroyed_accounts: B256Set,
145    ) {
146        // Retrieve updated nodes from hash builder.
147        let (_, updated_nodes) = hash_builder.split();
148        self.account_nodes.extend(exclude_empty_from_pair(updated_nodes));
149
150        // Add deleted node paths.
151        self.removed_nodes.extend(exclude_empty(removed_keys));
152
153        // Add deleted storage tries for destroyed accounts.
154        for destroyed in destroyed_accounts {
155            self.storage_tries.entry(destroyed).or_default().set_deleted(true);
156        }
157    }
158
159    /// Converts trie updates into [`TrieUpdatesSorted`].
160    pub fn into_sorted(mut self) -> TrieUpdatesSorted {
161        let mut account_nodes = self
162            .account_nodes
163            .drain()
164            .map(|(path, node)| {
165                // Updated nodes take precedence over removed nodes.
166                self.removed_nodes.remove(&path);
167                (path, Some(node))
168            })
169            .collect::<Vec<_>>();
170
171        account_nodes.extend(self.removed_nodes.drain().map(|path| (path, None)));
172        account_nodes.sort_unstable_by_key(|a| a.0);
173
174        let storage_tries = self
175            .storage_tries
176            .drain()
177            .map(|(hashed_address, updates)| (hashed_address, updates.into_sorted()))
178            .collect();
179        TrieUpdatesSorted { account_nodes, storage_tries }
180    }
181
182    /// Creates a sorted copy without consuming self.
183    /// More efficient than `.clone().into_sorted()` as it avoids cloning `HashMap` metadata.
184    pub fn clone_into_sorted(&self) -> TrieUpdatesSorted {
185        let mut account_nodes = self
186            .account_nodes
187            .iter()
188            .map(|(path, node)| (*path, Some(node.clone())))
189            .collect::<Vec<_>>();
190
191        // Add removed nodes that aren't already updated (updated nodes take precedence)
192        account_nodes.extend(
193            self.removed_nodes
194                .iter()
195                .filter(|path| !self.account_nodes.contains_key(*path))
196                .map(|path| (*path, None)),
197        );
198        account_nodes.sort_unstable_by_key(|a| a.0);
199
200        let storage_tries = self
201            .storage_tries
202            .iter()
203            .map(|(&hashed_address, updates)| (hashed_address, updates.clone_into_sorted()))
204            .collect();
205        TrieUpdatesSorted { account_nodes, storage_tries }
206    }
207
208    /// Converts trie updates into [`TrieUpdatesSortedRef`].
209    pub fn into_sorted_ref(&self) -> TrieUpdatesSortedRef<'_> {
210        let mut account_nodes = self.account_nodes.iter().collect::<Vec<_>>();
211        account_nodes.sort_unstable_by(|a, b| a.0.cmp(b.0));
212
213        TrieUpdatesSortedRef {
214            removed_nodes: self.removed_nodes.iter().collect::<BTreeSet<_>>(),
215            account_nodes,
216            storage_tries: self
217                .storage_tries
218                .iter()
219                .map(|m| (*m.0, m.1.into_sorted_ref()))
220                .collect(),
221        }
222    }
223
224    /// Clears the nodes and storage trie maps in this `TrieUpdates`.
225    pub fn clear(&mut self) {
226        self.account_nodes.clear();
227        self.removed_nodes.clear();
228        self.storage_tries.clear();
229    }
230}
231
232/// Trie updates for storage trie of a single account.
233#[derive(PartialEq, Eq, Clone, Default, Debug)]
234#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
235pub struct StorageTrieUpdates {
236    /// Flag indicating whether the trie was deleted.
237    pub is_deleted: bool,
238    /// Collection of updated storage trie nodes.
239    #[cfg_attr(any(test, feature = "serde"), serde(with = "serde_nibbles_map"))]
240    pub storage_nodes: HashMap<Nibbles, BranchNodeCompact>,
241    /// Collection of removed storage trie nodes.
242    #[cfg_attr(any(test, feature = "serde"), serde(with = "serde_nibbles_set"))]
243    pub removed_nodes: HashSet<Nibbles>,
244}
245
246#[cfg(feature = "test-utils")]
247impl StorageTrieUpdates {
248    /// Creates a new storage trie updates that are not marked as deleted.
249    pub fn new(updates: impl IntoIterator<Item = (Nibbles, BranchNodeCompact)>) -> Self {
250        Self { storage_nodes: exclude_empty_from_pair(updates).collect(), ..Default::default() }
251    }
252}
253
254impl StorageTrieUpdates {
255    /// Returns empty storage trie updates with `deleted` set to `true`.
256    pub fn deleted() -> Self {
257        Self {
258            is_deleted: true,
259            storage_nodes: HashMap::default(),
260            removed_nodes: HashSet::default(),
261        }
262    }
263
264    /// Returns the length of updated nodes.
265    pub fn len(&self) -> usize {
266        (self.is_deleted as usize) + self.storage_nodes.len() + self.removed_nodes.len()
267    }
268
269    /// Returns `true` if the trie was deleted.
270    pub const fn is_deleted(&self) -> bool {
271        self.is_deleted
272    }
273
274    /// Returns reference to updated storage nodes.
275    pub const fn storage_nodes_ref(&self) -> &HashMap<Nibbles, BranchNodeCompact> {
276        &self.storage_nodes
277    }
278
279    /// Returns reference to removed storage nodes.
280    pub const fn removed_nodes_ref(&self) -> &HashSet<Nibbles> {
281        &self.removed_nodes
282    }
283
284    /// Returns `true` if storage updates are empty.
285    pub fn is_empty(&self) -> bool {
286        !self.is_deleted && self.storage_nodes.is_empty() && self.removed_nodes.is_empty()
287    }
288
289    /// Sets `deleted` flag on the storage trie.
290    pub const fn set_deleted(&mut self, deleted: bool) {
291        self.is_deleted = deleted;
292    }
293
294    /// Extends storage trie updates.
295    pub fn extend(&mut self, other: Self) {
296        self.extend_common(&other);
297        self.storage_nodes.extend(exclude_empty_from_pair(other.storage_nodes));
298        self.removed_nodes.extend(exclude_empty(other.removed_nodes));
299    }
300
301    /// Extends storage trie updates.
302    ///
303    /// Slightly less efficient than [`Self::extend`], but preferred to `extend(other.clone())`.
304    pub fn extend_ref(&mut self, other: &Self) {
305        self.extend_common(other);
306        self.storage_nodes.extend(exclude_empty_from_pair(
307            other.storage_nodes.iter().map(|(k, v)| (*k, v.clone())),
308        ));
309        self.removed_nodes.extend(exclude_empty(other.removed_nodes.iter().copied()));
310    }
311
312    fn extend_common(&mut self, other: &Self) {
313        if other.is_deleted {
314            self.storage_nodes.clear();
315            self.removed_nodes.clear();
316        }
317        self.is_deleted |= other.is_deleted;
318        self.storage_nodes.retain(|nibbles, _| !other.removed_nodes.contains(nibbles));
319    }
320
321    /// Extend storage trie updates with sorted data, converting directly into the unsorted
322    /// `HashMap` representation. This is more efficient than first converting to
323    /// `StorageTrieUpdates` and then extending, as it avoids creating intermediate `HashMap`
324    /// allocations.
325    ///
326    /// This is invoked from [`TrieUpdates::extend_from_sorted`] for each account.
327    pub fn extend_from_sorted(&mut self, sorted: &StorageTrieUpdatesSorted) {
328        if sorted.is_deleted {
329            self.storage_nodes.clear();
330            self.removed_nodes.clear();
331        }
332        self.is_deleted |= sorted.is_deleted;
333
334        // Reserve capacity for storage nodes
335        let new_nodes_count = sorted.storage_nodes.len();
336        self.storage_nodes.reserve(new_nodes_count);
337
338        // Remove nodes marked as removed and insert new nodes
339        for (nibbles, maybe_node) in &sorted.storage_nodes {
340            if nibbles.is_empty() {
341                continue;
342            }
343            if let Some(node) = maybe_node {
344                self.removed_nodes.remove(nibbles);
345                self.storage_nodes.insert(*nibbles, node.clone());
346            } else {
347                self.storage_nodes.remove(nibbles);
348                self.removed_nodes.insert(*nibbles);
349            }
350        }
351    }
352
353    /// Finalize storage trie updates for by taking updates from walker and hash builder.
354    pub fn finalize(&mut self, hash_builder: HashBuilder, removed_keys: HashSet<Nibbles>) {
355        // Retrieve updated nodes from hash builder.
356        let (_, updated_nodes) = hash_builder.split();
357        self.storage_nodes.extend(exclude_empty_from_pair(updated_nodes));
358
359        // Add deleted node paths.
360        self.removed_nodes.extend(exclude_empty(removed_keys));
361    }
362
363    /// Convert storage trie updates into [`StorageTrieUpdatesSorted`].
364    pub fn into_sorted(mut self) -> StorageTrieUpdatesSorted {
365        let mut storage_nodes = self
366            .storage_nodes
367            .into_iter()
368            .map(|(path, node)| {
369                // Updated nodes take precedence over removed nodes.
370                self.removed_nodes.remove(&path);
371                (path, Some(node))
372            })
373            .collect::<Vec<_>>();
374
375        storage_nodes.extend(self.removed_nodes.into_iter().map(|path| (path, None)));
376        storage_nodes.sort_unstable_by_key(|a| a.0);
377
378        StorageTrieUpdatesSorted { is_deleted: self.is_deleted, storage_nodes }
379    }
380
381    /// Creates a sorted copy without consuming self.
382    /// More efficient than `.clone().into_sorted()` as it avoids cloning `HashMap` metadata.
383    pub fn clone_into_sorted(&self) -> StorageTrieUpdatesSorted {
384        let mut storage_nodes = self
385            .storage_nodes
386            .iter()
387            .map(|(path, node)| (*path, Some(node.clone())))
388            .collect::<Vec<_>>();
389
390        // Add removed nodes that aren't already updated (updated nodes take precedence)
391        storage_nodes.extend(
392            self.removed_nodes
393                .iter()
394                .filter(|path| !self.storage_nodes.contains_key(*path))
395                .map(|path| (*path, None)),
396        );
397        storage_nodes.sort_unstable_by_key(|a| a.0);
398
399        StorageTrieUpdatesSorted { is_deleted: self.is_deleted, storage_nodes }
400    }
401
402    /// Convert storage trie updates into [`StorageTrieUpdatesSortedRef`].
403    pub fn into_sorted_ref(&self) -> StorageTrieUpdatesSortedRef<'_> {
404        StorageTrieUpdatesSortedRef {
405            is_deleted: self.is_deleted,
406            removed_nodes: self.removed_nodes.iter().collect::<BTreeSet<_>>(),
407            storage_nodes: self.storage_nodes.iter().collect::<BTreeMap<_, _>>(),
408        }
409    }
410}
411
412/// Serializes and deserializes any [`HashSet`] that includes [`Nibbles`] elements, by using the
413/// hex-encoded packed representation.
414///
415/// This also sorts the set before serializing.
416#[cfg(any(test, feature = "serde"))]
417mod serde_nibbles_set {
418    use crate::Nibbles;
419    use alloc::{
420        string::{String, ToString},
421        vec::Vec,
422    };
423    use alloy_primitives::map::HashSet;
424    use serde::{de::Error, Deserialize, Deserializer, Serialize, Serializer};
425
426    pub(super) fn serialize<S>(map: &HashSet<Nibbles>, serializer: S) -> Result<S::Ok, S::Error>
427    where
428        S: Serializer,
429    {
430        let mut storage_nodes =
431            map.iter().map(|elem| alloy_primitives::hex::encode(elem.pack())).collect::<Vec<_>>();
432        storage_nodes.sort_unstable();
433        storage_nodes.serialize(serializer)
434    }
435
436    pub(super) fn deserialize<'de, D>(deserializer: D) -> Result<HashSet<Nibbles>, D::Error>
437    where
438        D: Deserializer<'de>,
439    {
440        Vec::<String>::deserialize(deserializer)?
441            .into_iter()
442            .map(|node| {
443                Ok(Nibbles::unpack(
444                    alloy_primitives::hex::decode(node)
445                        .map_err(|err| D::Error::custom(err.to_string()))?,
446                ))
447            })
448            .collect::<Result<HashSet<_>, _>>()
449    }
450}
451
452/// Serializes and deserializes any [`HashMap`] that uses [`Nibbles`] as keys, by using the
453/// hex-encoded packed representation.
454///
455/// This also sorts the map's keys before encoding and serializing.
456#[cfg(any(test, feature = "serde"))]
457mod serde_nibbles_map {
458    use crate::Nibbles;
459    use alloc::{
460        string::{String, ToString},
461        vec::Vec,
462    };
463    use alloy_primitives::{hex, map::HashMap};
464    use core::marker::PhantomData;
465    use serde::{
466        de::{Error, MapAccess, Visitor},
467        ser::SerializeMap,
468        Deserialize, Deserializer, Serialize, Serializer,
469    };
470
471    pub(super) fn serialize<S, T>(
472        map: &HashMap<Nibbles, T>,
473        serializer: S,
474    ) -> Result<S::Ok, S::Error>
475    where
476        S: Serializer,
477        T: Serialize,
478    {
479        let mut map_serializer = serializer.serialize_map(Some(map.len()))?;
480        let mut storage_nodes = Vec::from_iter(map);
481        storage_nodes.sort_unstable_by_key(|node| node.0);
482        for (k, v) in storage_nodes {
483            // pack, then hex encode the Nibbles
484            let packed = alloy_primitives::hex::encode(k.pack());
485            map_serializer.serialize_entry(&packed, &v)?;
486        }
487        map_serializer.end()
488    }
489
490    pub(super) fn deserialize<'de, D, T>(deserializer: D) -> Result<HashMap<Nibbles, T>, D::Error>
491    where
492        D: Deserializer<'de>,
493        T: Deserialize<'de>,
494    {
495        struct NibblesMapVisitor<T> {
496            marker: PhantomData<T>,
497        }
498
499        impl<'de, T> Visitor<'de> for NibblesMapVisitor<T>
500        where
501            T: Deserialize<'de>,
502        {
503            type Value = HashMap<Nibbles, T>;
504
505            fn expecting(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
506                formatter.write_str("a map with hex-encoded Nibbles keys")
507            }
508
509            fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
510            where
511                A: MapAccess<'de>,
512            {
513                let mut result = HashMap::with_capacity_and_hasher(
514                    map.size_hint().unwrap_or(0),
515                    Default::default(),
516                );
517
518                while let Some((key, value)) = map.next_entry::<String, T>()? {
519                    let decoded_key =
520                        hex::decode(&key).map_err(|err| Error::custom(err.to_string()))?;
521
522                    let nibbles = Nibbles::unpack(&decoded_key);
523
524                    result.insert(nibbles, value);
525                }
526
527                Ok(result)
528            }
529        }
530
531        deserializer.deserialize_map(NibblesMapVisitor { marker: PhantomData })
532    }
533}
534
535/// Sorted trie updates reference used for serializing trie to file.
536#[derive(PartialEq, Eq, Clone, Default, Debug)]
537#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize))]
538pub struct TrieUpdatesSortedRef<'a> {
539    /// Sorted collection of updated state nodes with corresponding paths.
540    pub account_nodes: Vec<(&'a Nibbles, &'a BranchNodeCompact)>,
541    /// The set of removed state node keys.
542    pub removed_nodes: BTreeSet<&'a Nibbles>,
543    /// Storage tries stored by hashed address of the account the trie belongs to.
544    pub storage_tries: BTreeMap<FixedBytes<32>, StorageTrieUpdatesSortedRef<'a>>,
545}
546
547/// Sorted trie updates used for lookups and insertions.
548#[derive(PartialEq, Eq, Clone, Default, Debug)]
549#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
550pub struct TrieUpdatesSorted {
551    /// Sorted collection of updated state nodes with corresponding paths. None indicates that a
552    /// node was removed.
553    account_nodes: Vec<(Nibbles, Option<BranchNodeCompact>)>,
554    /// Storage tries stored by hashed address of the account the trie belongs to.
555    storage_tries: B256Map<StorageTrieUpdatesSorted>,
556}
557
558impl TrieUpdatesSorted {
559    /// Creates a new `TrieUpdatesSorted` with the given account nodes and storage tries.
560    ///
561    /// # Panics
562    ///
563    /// In debug mode, panics if `account_nodes` is not sorted by the `Nibbles` key,
564    /// or if any storage trie's `storage_nodes` is not sorted by its `Nibbles` key.
565    pub fn new(
566        account_nodes: Vec<(Nibbles, Option<BranchNodeCompact>)>,
567        storage_tries: B256Map<StorageTrieUpdatesSorted>,
568    ) -> Self {
569        debug_assert!(
570            account_nodes.is_sorted_by_key(|item| &item.0),
571            "account_nodes must be sorted by Nibbles key"
572        );
573        debug_assert!(
574            storage_tries.values().all(|storage_trie| {
575                storage_trie.storage_nodes.is_sorted_by_key(|item| &item.0)
576            }),
577            "all storage_nodes in storage_tries must be sorted by Nibbles key"
578        );
579        Self { account_nodes, storage_tries }
580    }
581
582    /// Returns `true` if the updates are empty.
583    pub fn is_empty(&self) -> bool {
584        self.account_nodes.is_empty() && self.storage_tries.is_empty()
585    }
586
587    /// Returns reference to updated account nodes.
588    pub fn account_nodes_ref(&self) -> &[(Nibbles, Option<BranchNodeCompact>)] {
589        &self.account_nodes
590    }
591
592    /// Returns reference to updated storage tries.
593    pub const fn storage_tries_ref(&self) -> &B256Map<StorageTrieUpdatesSorted> {
594        &self.storage_tries
595    }
596
597    /// Returns the total number of updates including account nodes and all storage updates.
598    pub fn total_len(&self) -> usize {
599        self.account_nodes.len() +
600            self.storage_tries.values().map(|storage| storage.len()).sum::<usize>()
601    }
602
603    /// Extends the trie updates with another set of sorted updates.
604    ///
605    /// This merges the account nodes and storage tries from `other` into `self`.
606    /// Account nodes are merged and re-sorted, with `other`'s values taking precedence
607    /// for duplicate keys.
608    ///
609    /// Sorts the account nodes after extending. Sorts the storage tries after extending, for each
610    /// storage trie.
611    pub fn extend_ref_and_sort(&mut self, other: &Self) {
612        // Extend account nodes
613        extend_sorted_vec(&mut self.account_nodes, &other.account_nodes);
614
615        // Merge storage tries
616        for (hashed_address, storage_trie) in &other.storage_tries {
617            self.storage_tries
618                .entry(*hashed_address)
619                .and_modify(|existing| existing.extend_ref(storage_trie))
620                .or_insert_with(|| storage_trie.clone());
621        }
622    }
623
624    /// Clears all account nodes and storage tries.
625    pub fn clear(&mut self) {
626        self.account_nodes.clear();
627        self.storage_tries.clear();
628    }
629
630    /// Batch-merge sorted trie updates. Iterator yields **newest to oldest**.
631    ///
632    /// For small batches, uses `extend_ref_and_sort` loop.
633    /// For large batches, uses k-way merge for O(n log k) complexity.
634    pub fn merge_batch<T: AsRef<Self> + From<Self>>(iter: impl IntoIterator<Item = T>) -> T {
635        let items: alloc::vec::Vec<_> = iter.into_iter().collect();
636        match items.len() {
637            0 => Self::default().into(),
638            1 => items.into_iter().next().expect("len == 1"),
639            _ => Self::merge_slice(&items).into(),
640        }
641    }
642
643    /// Batch-merge sorted trie updates from a slice. Slice is **newest to oldest**.
644    ///
645    /// This variant takes a slice reference directly, avoiding iterator collection overhead.
646    /// For small batches, uses `extend_ref_and_sort` loop.
647    /// For large batches, uses k-way merge for O(n log k) complexity.
648    pub fn merge_slice<T: AsRef<Self>>(items: &[T]) -> Self {
649        const THRESHOLD: usize = 30;
650
651        let k = items.len();
652
653        if k == 0 {
654            return Self::default();
655        }
656        if k == 1 {
657            return items[0].as_ref().clone();
658        }
659
660        if k < THRESHOLD {
661            // Small k: extend loop, oldest-to-newest so newer overrides older.
662            let mut iter = items.iter().rev();
663            let mut acc = iter.next().expect("k > 0").as_ref().clone();
664            for next in iter {
665                acc.extend_ref_and_sort(next.as_ref());
666            }
667            return acc;
668        }
669
670        // Large k: k-way merge.
671        let account_nodes =
672            kway_merge_sorted(items.iter().map(|i| i.as_ref().account_nodes.as_slice()));
673
674        struct StorageAcc<'a> {
675            is_deleted: bool,
676            sealed: bool,
677            slices: Vec<&'a [(Nibbles, Option<BranchNodeCompact>)]>,
678        }
679
680        let mut acc: B256Map<StorageAcc<'_>> = B256Map::default();
681
682        for item in items {
683            for (addr, storage) in &item.as_ref().storage_tries {
684                let entry = acc.entry(*addr).or_insert_with(|| StorageAcc {
685                    is_deleted: false,
686                    sealed: false,
687                    slices: Vec::new(),
688                });
689
690                if entry.sealed {
691                    continue;
692                }
693
694                entry.slices.push(storage.storage_nodes.as_slice());
695
696                if storage.is_deleted {
697                    entry.is_deleted = true;
698                    entry.sealed = true;
699                }
700            }
701        }
702
703        let storage_tries = acc
704            .into_iter()
705            .map(|(addr, entry)| {
706                let storage_nodes = kway_merge_sorted(entry.slices);
707                (addr, StorageTrieUpdatesSorted { is_deleted: entry.is_deleted, storage_nodes })
708            })
709            .collect();
710
711        Self { account_nodes, storage_tries }
712    }
713}
714
715impl AsRef<Self> for TrieUpdatesSorted {
716    fn as_ref(&self) -> &Self {
717        self
718    }
719}
720
721impl From<TrieUpdatesSorted> for TrieUpdates {
722    fn from(sorted: TrieUpdatesSorted) -> Self {
723        let mut account_nodes = HashMap::default();
724        let mut removed_nodes = HashSet::default();
725
726        for (nibbles, node) in sorted.account_nodes {
727            if let Some(node) = node {
728                account_nodes.insert(nibbles, node);
729            } else {
730                removed_nodes.insert(nibbles);
731            }
732        }
733
734        let storage_tries = sorted
735            .storage_tries
736            .into_iter()
737            .map(|(address, storage)| (address, storage.into()))
738            .collect();
739
740        Self { account_nodes, removed_nodes, storage_tries }
741    }
742}
743
744/// Sorted storage trie updates reference used for serializing to file.
745#[derive(PartialEq, Eq, Clone, Default, Debug)]
746#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize))]
747pub struct StorageTrieUpdatesSortedRef<'a> {
748    /// Flag indicating whether the trie has been deleted/wiped.
749    pub is_deleted: bool,
750    /// Sorted collection of updated storage nodes with corresponding paths.
751    pub storage_nodes: BTreeMap<&'a Nibbles, &'a BranchNodeCompact>,
752    /// The set of removed storage node keys.
753    pub removed_nodes: BTreeSet<&'a Nibbles>,
754}
755
756/// Sorted trie updates used for lookups and insertions.
757#[derive(PartialEq, Eq, Clone, Default, Debug)]
758#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
759pub struct StorageTrieUpdatesSorted {
760    /// Flag indicating whether the trie has been deleted/wiped.
761    pub is_deleted: bool,
762    /// Sorted collection of updated storage nodes with corresponding paths. None indicates a node
763    /// is removed.
764    pub storage_nodes: Vec<(Nibbles, Option<BranchNodeCompact>)>,
765}
766
767impl StorageTrieUpdatesSorted {
768    /// Returns `true` if the trie was deleted.
769    pub const fn is_deleted(&self) -> bool {
770        self.is_deleted
771    }
772
773    /// Returns reference to updated storage nodes.
774    pub fn storage_nodes_ref(&self) -> &[(Nibbles, Option<BranchNodeCompact>)] {
775        &self.storage_nodes
776    }
777
778    /// Returns the total number of storage node updates.
779    pub const fn len(&self) -> usize {
780        self.storage_nodes.len()
781    }
782
783    /// Returns `true` if there are no storage node updates.
784    pub const fn is_empty(&self) -> bool {
785        self.storage_nodes.is_empty()
786    }
787
788    /// Extends the storage trie updates with another set of sorted updates.
789    ///
790    /// If `other` is marked as deleted, this will be marked as deleted and all nodes cleared.
791    /// Otherwise, nodes are merged with `other`'s values taking precedence for duplicates.
792    pub fn extend_ref(&mut self, other: &Self) {
793        if other.is_deleted {
794            self.is_deleted = true;
795            self.storage_nodes.clear();
796            self.storage_nodes.extend(other.storage_nodes.iter().cloned());
797            return;
798        }
799
800        // Extend storage nodes
801        extend_sorted_vec(&mut self.storage_nodes, &other.storage_nodes);
802        self.is_deleted = self.is_deleted || other.is_deleted;
803    }
804
805    /// Batch-merge sorted storage trie updates. Iterator yields **newest to oldest**.
806    /// If any update is deleted, older data is discarded.
807    pub fn merge_batch<'a>(updates: impl IntoIterator<Item = &'a Self>) -> Self {
808        let updates: Vec<_> = updates.into_iter().collect();
809        if updates.is_empty() {
810            return Self::default();
811        }
812
813        // Discard updates older than the first deletion since the trie was wiped at that point.
814        let del_idx = updates.iter().position(|u| u.is_deleted);
815        let relevant = del_idx.map_or(&updates[..], |idx| &updates[..=idx]);
816        let storage_nodes = kway_merge_sorted(relevant.iter().map(|u| u.storage_nodes.as_slice()));
817
818        Self { is_deleted: del_idx.is_some(), storage_nodes }
819    }
820}
821
822/// Excludes empty nibbles from the given iterator.
823fn exclude_empty(iter: impl IntoIterator<Item = Nibbles>) -> impl Iterator<Item = Nibbles> {
824    iter.into_iter().filter(|n| !n.is_empty())
825}
826
827/// Excludes empty nibbles from the given iterator of pairs where the nibbles are the key.
828fn exclude_empty_from_pair<V>(
829    iter: impl IntoIterator<Item = (Nibbles, V)>,
830) -> impl Iterator<Item = (Nibbles, V)> {
831    iter.into_iter().filter(|(n, _)| !n.is_empty())
832}
833
834impl From<StorageTrieUpdatesSorted> for StorageTrieUpdates {
835    fn from(sorted: StorageTrieUpdatesSorted) -> Self {
836        let mut storage_nodes = HashMap::default();
837        let mut removed_nodes = HashSet::default();
838
839        for (nibbles, node) in sorted.storage_nodes {
840            if let Some(node) = node {
841                storage_nodes.insert(nibbles, node);
842            } else {
843                removed_nodes.insert(nibbles);
844            }
845        }
846
847        Self { is_deleted: sorted.is_deleted, storage_nodes, removed_nodes }
848    }
849}
850
851#[cfg(test)]
852mod tests {
853    use super::*;
854    use alloy_primitives::B256;
855
856    #[test]
857    fn test_trie_updates_sorted_extend_ref() {
858        // Test extending with empty updates
859        let mut updates1 = TrieUpdatesSorted::default();
860        let updates2 = TrieUpdatesSorted::default();
861        updates1.extend_ref_and_sort(&updates2);
862        assert_eq!(updates1.account_nodes.len(), 0);
863        assert_eq!(updates1.storage_tries.len(), 0);
864
865        // Test extending account nodes
866        let mut updates1 = TrieUpdatesSorted {
867            account_nodes: vec![
868                (Nibbles::from_nibbles_unchecked([0x01]), Some(BranchNodeCompact::default())),
869                (Nibbles::from_nibbles_unchecked([0x03]), None),
870            ],
871            storage_tries: B256Map::default(),
872        };
873        let updates2 = TrieUpdatesSorted {
874            account_nodes: vec![
875                (Nibbles::from_nibbles_unchecked([0x02]), Some(BranchNodeCompact::default())),
876                (Nibbles::from_nibbles_unchecked([0x03]), Some(BranchNodeCompact::default())), /* Override */
877            ],
878            storage_tries: B256Map::default(),
879        };
880        updates1.extend_ref_and_sort(&updates2);
881        assert_eq!(updates1.account_nodes.len(), 3);
882        // Should be sorted: 0x01, 0x02, 0x03
883        assert_eq!(updates1.account_nodes[0].0, Nibbles::from_nibbles_unchecked([0x01]));
884        assert_eq!(updates1.account_nodes[1].0, Nibbles::from_nibbles_unchecked([0x02]));
885        assert_eq!(updates1.account_nodes[2].0, Nibbles::from_nibbles_unchecked([0x03]));
886        // 0x03 should have Some value from updates2 (override)
887        assert!(updates1.account_nodes[2].1.is_some());
888
889        // Test extending storage tries
890        let storage_trie1 = StorageTrieUpdatesSorted {
891            is_deleted: false,
892            storage_nodes: vec![(
893                Nibbles::from_nibbles_unchecked([0x0a]),
894                Some(BranchNodeCompact::default()),
895            )],
896        };
897        let storage_trie2 = StorageTrieUpdatesSorted {
898            is_deleted: false,
899            storage_nodes: vec![(Nibbles::from_nibbles_unchecked([0x0b]), None)],
900        };
901
902        let hashed_address1 = B256::from([1; 32]);
903        let hashed_address2 = B256::from([2; 32]);
904
905        let mut updates1 = TrieUpdatesSorted {
906            account_nodes: vec![],
907            storage_tries: B256Map::from_iter([(hashed_address1, storage_trie1.clone())]),
908        };
909        let updates2 = TrieUpdatesSorted {
910            account_nodes: vec![],
911            storage_tries: B256Map::from_iter([
912                (hashed_address1, storage_trie2),
913                (hashed_address2, storage_trie1),
914            ]),
915        };
916        updates1.extend_ref_and_sort(&updates2);
917        assert_eq!(updates1.storage_tries.len(), 2);
918        assert!(updates1.storage_tries.contains_key(&hashed_address1));
919        assert!(updates1.storage_tries.contains_key(&hashed_address2));
920        // Check that storage trie for hashed_address1 was extended
921        let merged_storage = &updates1.storage_tries[&hashed_address1];
922        assert_eq!(merged_storage.storage_nodes.len(), 2);
923    }
924
925    #[test]
926    fn test_storage_trie_updates_sorted_extend_ref_deleted() {
927        // Test case 1: Extending with a deleted storage trie that has nodes
928        let mut storage1 = StorageTrieUpdatesSorted {
929            is_deleted: false,
930            storage_nodes: vec![
931                (Nibbles::from_nibbles_unchecked([0x01]), Some(BranchNodeCompact::default())),
932                (Nibbles::from_nibbles_unchecked([0x02]), None),
933            ],
934        };
935
936        let storage2 = StorageTrieUpdatesSorted {
937            is_deleted: true,
938            storage_nodes: vec![
939                (Nibbles::from_nibbles_unchecked([0x03]), Some(BranchNodeCompact::default())),
940                (Nibbles::from_nibbles_unchecked([0x04]), None),
941            ],
942        };
943
944        storage1.extend_ref(&storage2);
945
946        // Should be marked as deleted
947        assert!(storage1.is_deleted);
948        // Original nodes should be cleared, but other's nodes should be added
949        assert_eq!(storage1.storage_nodes.len(), 2);
950        assert_eq!(storage1.storage_nodes[0].0, Nibbles::from_nibbles_unchecked([0x03]));
951        assert_eq!(storage1.storage_nodes[1].0, Nibbles::from_nibbles_unchecked([0x04]));
952
953        // Test case 2: Extending a deleted storage trie with more nodes
954        let mut storage3 = StorageTrieUpdatesSorted {
955            is_deleted: true,
956            storage_nodes: vec![(
957                Nibbles::from_nibbles_unchecked([0x05]),
958                Some(BranchNodeCompact::default()),
959            )],
960        };
961
962        let storage4 = StorageTrieUpdatesSorted {
963            is_deleted: true,
964            storage_nodes: vec![
965                (Nibbles::from_nibbles_unchecked([0x06]), Some(BranchNodeCompact::default())),
966                (Nibbles::from_nibbles_unchecked([0x07]), None),
967            ],
968        };
969
970        storage3.extend_ref(&storage4);
971
972        // Should remain deleted
973        assert!(storage3.is_deleted);
974        // Should have nodes from other (original cleared then extended)
975        assert_eq!(storage3.storage_nodes.len(), 2);
976        assert_eq!(storage3.storage_nodes[0].0, Nibbles::from_nibbles_unchecked([0x06]));
977        assert_eq!(storage3.storage_nodes[1].0, Nibbles::from_nibbles_unchecked([0x07]));
978    }
979
980    /// Test extending with storage tries adds both nodes and removed nodes correctly
981    #[test]
982    fn test_trie_updates_extend_from_sorted_with_storage_tries() {
983        let hashed_address = B256::from([1; 32]);
984
985        let mut updates = TrieUpdates::default();
986
987        let storage_trie = StorageTrieUpdatesSorted {
988            is_deleted: false,
989            storage_nodes: vec![
990                (Nibbles::from_nibbles_unchecked([0x0a]), Some(BranchNodeCompact::default())),
991                (Nibbles::from_nibbles_unchecked([0x0b]), None),
992            ],
993        };
994
995        let sorted = TrieUpdatesSorted {
996            account_nodes: vec![],
997            storage_tries: B256Map::from_iter([(hashed_address, storage_trie)]),
998        };
999
1000        updates.extend_from_sorted(&sorted);
1001
1002        assert_eq!(updates.storage_tries.len(), 1);
1003        let storage = updates.storage_tries.get(&hashed_address).unwrap();
1004        assert!(!storage.is_deleted);
1005        assert_eq!(storage.storage_nodes.len(), 1);
1006        assert!(storage.removed_nodes.contains(&Nibbles::from_nibbles_unchecked([0x0b])));
1007    }
1008
1009    /// Test deleted=true clears old storage nodes before adding new ones (critical edge case)
1010    #[test]
1011    fn test_trie_updates_extend_from_sorted_with_deleted_storage() {
1012        let hashed_address = B256::from([1; 32]);
1013
1014        let mut updates = TrieUpdates::default();
1015        updates.storage_tries.insert(
1016            hashed_address,
1017            StorageTrieUpdates {
1018                is_deleted: false,
1019                storage_nodes: HashMap::from_iter([(
1020                    Nibbles::from_nibbles_unchecked([0x01]),
1021                    BranchNodeCompact::default(),
1022                )]),
1023                removed_nodes: Default::default(),
1024            },
1025        );
1026
1027        let storage_trie = StorageTrieUpdatesSorted {
1028            is_deleted: true,
1029            storage_nodes: vec![(
1030                Nibbles::from_nibbles_unchecked([0x0a]),
1031                Some(BranchNodeCompact::default()),
1032            )],
1033        };
1034
1035        let sorted = TrieUpdatesSorted {
1036            account_nodes: vec![],
1037            storage_tries: B256Map::from_iter([(hashed_address, storage_trie)]),
1038        };
1039
1040        updates.extend_from_sorted(&sorted);
1041
1042        let storage = updates.storage_tries.get(&hashed_address).unwrap();
1043        assert!(storage.is_deleted);
1044        // After deletion, old nodes should be cleared
1045        assert_eq!(storage.storage_nodes.len(), 1);
1046        assert!(storage.storage_nodes.contains_key(&Nibbles::from_nibbles_unchecked([0x0a])));
1047    }
1048
1049    /// Test non-deleted storage merges nodes and tracks removed nodes
1050    #[test]
1051    fn test_storage_trie_updates_extend_from_sorted_non_deleted() {
1052        let mut storage = StorageTrieUpdates {
1053            is_deleted: false,
1054            storage_nodes: HashMap::from_iter([(
1055                Nibbles::from_nibbles_unchecked([0x01]),
1056                BranchNodeCompact::default(),
1057            )]),
1058            removed_nodes: Default::default(),
1059        };
1060
1061        let sorted = StorageTrieUpdatesSorted {
1062            is_deleted: false,
1063            storage_nodes: vec![
1064                (Nibbles::from_nibbles_unchecked([0x02]), Some(BranchNodeCompact::default())),
1065                (Nibbles::from_nibbles_unchecked([0x03]), None),
1066            ],
1067        };
1068
1069        storage.extend_from_sorted(&sorted);
1070
1071        assert!(!storage.is_deleted);
1072        assert_eq!(storage.storage_nodes.len(), 2);
1073        assert!(storage.removed_nodes.contains(&Nibbles::from_nibbles_unchecked([0x03])));
1074    }
1075
1076    /// Test deleted=true clears old nodes before extending (edge case)
1077    #[test]
1078    fn test_storage_trie_updates_extend_from_sorted_deleted() {
1079        let mut storage = StorageTrieUpdates {
1080            is_deleted: false,
1081            storage_nodes: HashMap::from_iter([(
1082                Nibbles::from_nibbles_unchecked([0x01]),
1083                BranchNodeCompact::default(),
1084            )]),
1085            removed_nodes: Default::default(),
1086        };
1087
1088        let sorted = StorageTrieUpdatesSorted {
1089            is_deleted: true,
1090            storage_nodes: vec![(
1091                Nibbles::from_nibbles_unchecked([0x0a]),
1092                Some(BranchNodeCompact::default()),
1093            )],
1094        };
1095
1096        storage.extend_from_sorted(&sorted);
1097
1098        assert!(storage.is_deleted);
1099        // Old nodes should be cleared when deleted
1100        assert_eq!(storage.storage_nodes.len(), 1);
1101        assert!(storage.storage_nodes.contains_key(&Nibbles::from_nibbles_unchecked([0x0a])));
1102    }
1103
1104    /// Test empty nibbles are filtered out during conversion (edge case bug)
1105    #[test]
1106    fn test_trie_updates_extend_from_sorted_filters_empty_nibbles() {
1107        let mut updates = TrieUpdates::default();
1108
1109        let sorted = TrieUpdatesSorted {
1110            account_nodes: vec![
1111                (Nibbles::default(), Some(BranchNodeCompact::default())), // Empty nibbles
1112                (Nibbles::from_nibbles_unchecked([0x01]), Some(BranchNodeCompact::default())),
1113            ],
1114            storage_tries: B256Map::default(),
1115        };
1116
1117        updates.extend_from_sorted(&sorted);
1118
1119        // Empty nibbles should be filtered out
1120        assert_eq!(updates.account_nodes.len(), 1);
1121        assert!(updates.account_nodes.contains_key(&Nibbles::from_nibbles_unchecked([0x01])));
1122        assert!(!updates.account_nodes.contains_key(&Nibbles::default()));
1123    }
1124}
1125
1126/// Bincode-compatible trie updates type serde implementations.
1127#[cfg(feature = "serde-bincode-compat")]
1128pub mod serde_bincode_compat {
1129    use crate::{BranchNodeCompact, Nibbles};
1130    use alloc::borrow::Cow;
1131    use alloy_primitives::map::{B256Map, HashMap, HashSet};
1132    use serde::{Deserialize, Deserializer, Serialize, Serializer};
1133    use serde_with::{DeserializeAs, SerializeAs};
1134
1135    /// Bincode-compatible [`super::TrieUpdates`] serde implementation.
1136    ///
1137    /// Intended to use with the [`serde_with::serde_as`] macro in the following way:
1138    /// ```rust
1139    /// use reth_trie_common::{serde_bincode_compat, updates::TrieUpdates};
1140    /// use serde::{Deserialize, Serialize};
1141    /// use serde_with::serde_as;
1142    ///
1143    /// #[serde_as]
1144    /// #[derive(Serialize, Deserialize)]
1145    /// struct Data {
1146    ///     #[serde_as(as = "serde_bincode_compat::updates::TrieUpdates")]
1147    ///     trie_updates: TrieUpdates,
1148    /// }
1149    /// ```
1150    #[derive(Debug, Serialize, Deserialize)]
1151    pub struct TrieUpdates<'a> {
1152        account_nodes: Cow<'a, HashMap<Nibbles, BranchNodeCompact>>,
1153        removed_nodes: Cow<'a, HashSet<Nibbles>>,
1154        storage_tries: B256Map<StorageTrieUpdates<'a>>,
1155    }
1156
1157    impl<'a> From<&'a super::TrieUpdates> for TrieUpdates<'a> {
1158        fn from(value: &'a super::TrieUpdates) -> Self {
1159            Self {
1160                account_nodes: Cow::Borrowed(&value.account_nodes),
1161                removed_nodes: Cow::Borrowed(&value.removed_nodes),
1162                storage_tries: value.storage_tries.iter().map(|(k, v)| (*k, v.into())).collect(),
1163            }
1164        }
1165    }
1166
1167    impl<'a> From<TrieUpdates<'a>> for super::TrieUpdates {
1168        fn from(value: TrieUpdates<'a>) -> Self {
1169            Self {
1170                account_nodes: value.account_nodes.into_owned(),
1171                removed_nodes: value.removed_nodes.into_owned(),
1172                storage_tries: value
1173                    .storage_tries
1174                    .into_iter()
1175                    .map(|(k, v)| (k, v.into()))
1176                    .collect(),
1177            }
1178        }
1179    }
1180
1181    impl SerializeAs<super::TrieUpdates> for TrieUpdates<'_> {
1182        fn serialize_as<S>(source: &super::TrieUpdates, serializer: S) -> Result<S::Ok, S::Error>
1183        where
1184            S: Serializer,
1185        {
1186            TrieUpdates::from(source).serialize(serializer)
1187        }
1188    }
1189
1190    impl<'de> DeserializeAs<'de, super::TrieUpdates> for TrieUpdates<'de> {
1191        fn deserialize_as<D>(deserializer: D) -> Result<super::TrieUpdates, D::Error>
1192        where
1193            D: Deserializer<'de>,
1194        {
1195            TrieUpdates::deserialize(deserializer).map(Into::into)
1196        }
1197    }
1198
1199    /// Bincode-compatible [`super::StorageTrieUpdates`] serde implementation.
1200    ///
1201    /// Intended to use with the [`serde_with::serde_as`] macro in the following way:
1202    /// ```rust
1203    /// use reth_trie_common::{serde_bincode_compat, updates::StorageTrieUpdates};
1204    /// use serde::{Deserialize, Serialize};
1205    /// use serde_with::serde_as;
1206    ///
1207    /// #[serde_as]
1208    /// #[derive(Serialize, Deserialize)]
1209    /// struct Data {
1210    ///     #[serde_as(as = "serde_bincode_compat::updates::StorageTrieUpdates")]
1211    ///     trie_updates: StorageTrieUpdates,
1212    /// }
1213    /// ```
1214    #[derive(Debug, Serialize, Deserialize)]
1215    pub struct StorageTrieUpdates<'a> {
1216        is_deleted: bool,
1217        storage_nodes: Cow<'a, HashMap<Nibbles, BranchNodeCompact>>,
1218        removed_nodes: Cow<'a, HashSet<Nibbles>>,
1219    }
1220
1221    impl<'a> From<&'a super::StorageTrieUpdates> for StorageTrieUpdates<'a> {
1222        fn from(value: &'a super::StorageTrieUpdates) -> Self {
1223            Self {
1224                is_deleted: value.is_deleted,
1225                storage_nodes: Cow::Borrowed(&value.storage_nodes),
1226                removed_nodes: Cow::Borrowed(&value.removed_nodes),
1227            }
1228        }
1229    }
1230
1231    impl<'a> From<StorageTrieUpdates<'a>> for super::StorageTrieUpdates {
1232        fn from(value: StorageTrieUpdates<'a>) -> Self {
1233            Self {
1234                is_deleted: value.is_deleted,
1235                storage_nodes: value.storage_nodes.into_owned(),
1236                removed_nodes: value.removed_nodes.into_owned(),
1237            }
1238        }
1239    }
1240
1241    impl SerializeAs<super::StorageTrieUpdates> for StorageTrieUpdates<'_> {
1242        fn serialize_as<S>(
1243            source: &super::StorageTrieUpdates,
1244            serializer: S,
1245        ) -> Result<S::Ok, S::Error>
1246        where
1247            S: Serializer,
1248        {
1249            StorageTrieUpdates::from(source).serialize(serializer)
1250        }
1251    }
1252
1253    impl<'de> DeserializeAs<'de, super::StorageTrieUpdates> for StorageTrieUpdates<'de> {
1254        fn deserialize_as<D>(deserializer: D) -> Result<super::StorageTrieUpdates, D::Error>
1255        where
1256            D: Deserializer<'de>,
1257        {
1258            StorageTrieUpdates::deserialize(deserializer).map(Into::into)
1259        }
1260    }
1261
1262    /// Bincode-compatible [`super::TrieUpdatesSorted`] serde implementation.
1263    ///
1264    /// Intended to use with the [`serde_with::serde_as`] macro in the following way:
1265    /// ```rust
1266    /// use reth_trie_common::{serde_bincode_compat, updates::TrieUpdatesSorted};
1267    /// use serde::{Deserialize, Serialize};
1268    /// use serde_with::serde_as;
1269    ///
1270    /// #[serde_as]
1271    /// #[derive(Serialize, Deserialize)]
1272    /// struct Data {
1273    ///     #[serde_as(as = "serde_bincode_compat::updates::TrieUpdatesSorted")]
1274    ///     trie_updates: TrieUpdatesSorted,
1275    /// }
1276    /// ```
1277    #[derive(Debug, Serialize, Deserialize)]
1278    pub struct TrieUpdatesSorted<'a> {
1279        account_nodes: Cow<'a, [(Nibbles, Option<BranchNodeCompact>)]>,
1280        storage_tries: B256Map<StorageTrieUpdatesSorted<'a>>,
1281    }
1282
1283    impl<'a> From<&'a super::TrieUpdatesSorted> for TrieUpdatesSorted<'a> {
1284        fn from(value: &'a super::TrieUpdatesSorted) -> Self {
1285            Self {
1286                account_nodes: Cow::Borrowed(&value.account_nodes),
1287                storage_tries: value.storage_tries.iter().map(|(k, v)| (*k, v.into())).collect(),
1288            }
1289        }
1290    }
1291
1292    impl<'a> From<TrieUpdatesSorted<'a>> for super::TrieUpdatesSorted {
1293        fn from(value: TrieUpdatesSorted<'a>) -> Self {
1294            Self {
1295                account_nodes: value.account_nodes.into_owned(),
1296                storage_tries: value
1297                    .storage_tries
1298                    .into_iter()
1299                    .map(|(k, v)| (k, v.into()))
1300                    .collect(),
1301            }
1302        }
1303    }
1304
1305    impl SerializeAs<super::TrieUpdatesSorted> for TrieUpdatesSorted<'_> {
1306        fn serialize_as<S>(
1307            source: &super::TrieUpdatesSorted,
1308            serializer: S,
1309        ) -> Result<S::Ok, S::Error>
1310        where
1311            S: Serializer,
1312        {
1313            TrieUpdatesSorted::from(source).serialize(serializer)
1314        }
1315    }
1316
1317    impl<'de> DeserializeAs<'de, super::TrieUpdatesSorted> for TrieUpdatesSorted<'de> {
1318        fn deserialize_as<D>(deserializer: D) -> Result<super::TrieUpdatesSorted, D::Error>
1319        where
1320            D: Deserializer<'de>,
1321        {
1322            TrieUpdatesSorted::deserialize(deserializer).map(Into::into)
1323        }
1324    }
1325
1326    /// Bincode-compatible [`super::StorageTrieUpdatesSorted`] serde implementation.
1327    ///
1328    /// Intended to use with the [`serde_with::serde_as`] macro in the following way:
1329    /// ```rust
1330    /// use reth_trie_common::{serde_bincode_compat, updates::StorageTrieUpdatesSorted};
1331    /// use serde::{Deserialize, Serialize};
1332    /// use serde_with::serde_as;
1333    ///
1334    /// #[serde_as]
1335    /// #[derive(Serialize, Deserialize)]
1336    /// struct Data {
1337    ///     #[serde_as(as = "serde_bincode_compat::updates::StorageTrieUpdatesSorted")]
1338    ///     trie_updates: StorageTrieUpdatesSorted,
1339    /// }
1340    /// ```
1341    #[derive(Debug, Serialize, Deserialize)]
1342    pub struct StorageTrieUpdatesSorted<'a> {
1343        is_deleted: bool,
1344        storage_nodes: Cow<'a, [(Nibbles, Option<BranchNodeCompact>)]>,
1345    }
1346
1347    impl<'a> From<&'a super::StorageTrieUpdatesSorted> for StorageTrieUpdatesSorted<'a> {
1348        fn from(value: &'a super::StorageTrieUpdatesSorted) -> Self {
1349            Self {
1350                is_deleted: value.is_deleted,
1351                storage_nodes: Cow::Borrowed(&value.storage_nodes),
1352            }
1353        }
1354    }
1355
1356    impl<'a> From<StorageTrieUpdatesSorted<'a>> for super::StorageTrieUpdatesSorted {
1357        fn from(value: StorageTrieUpdatesSorted<'a>) -> Self {
1358            Self { is_deleted: value.is_deleted, storage_nodes: value.storage_nodes.into_owned() }
1359        }
1360    }
1361
1362    impl SerializeAs<super::StorageTrieUpdatesSorted> for StorageTrieUpdatesSorted<'_> {
1363        fn serialize_as<S>(
1364            source: &super::StorageTrieUpdatesSorted,
1365            serializer: S,
1366        ) -> Result<S::Ok, S::Error>
1367        where
1368            S: Serializer,
1369        {
1370            StorageTrieUpdatesSorted::from(source).serialize(serializer)
1371        }
1372    }
1373
1374    impl<'de> DeserializeAs<'de, super::StorageTrieUpdatesSorted> for StorageTrieUpdatesSorted<'de> {
1375        fn deserialize_as<D>(deserializer: D) -> Result<super::StorageTrieUpdatesSorted, D::Error>
1376        where
1377            D: Deserializer<'de>,
1378        {
1379            StorageTrieUpdatesSorted::deserialize(deserializer).map(Into::into)
1380        }
1381    }
1382
1383    #[cfg(test)]
1384    mod tests {
1385        use crate::{
1386            serde_bincode_compat,
1387            updates::{
1388                StorageTrieUpdates, StorageTrieUpdatesSorted, TrieUpdates, TrieUpdatesSorted,
1389            },
1390            BranchNodeCompact, Nibbles,
1391        };
1392        use alloy_primitives::B256;
1393        use serde::{Deserialize, Serialize};
1394        use serde_with::serde_as;
1395
1396        #[test]
1397        fn test_trie_updates_bincode_roundtrip() {
1398            #[serde_as]
1399            #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
1400            struct Data {
1401                #[serde_as(as = "serde_bincode_compat::updates::TrieUpdates")]
1402                trie_updates: TrieUpdates,
1403            }
1404
1405            let mut data = Data { trie_updates: TrieUpdates::default() };
1406            let encoded = bincode::serialize(&data).unwrap();
1407            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1408            assert_eq!(decoded, data);
1409
1410            data.trie_updates
1411                .removed_nodes
1412                .insert(Nibbles::from_nibbles_unchecked([0x0b, 0x0e, 0x0e, 0x0f]));
1413            let encoded = bincode::serialize(&data).unwrap();
1414            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1415            assert_eq!(decoded, data);
1416
1417            data.trie_updates.account_nodes.insert(
1418                Nibbles::from_nibbles_unchecked([0x0d, 0x0e, 0x0a, 0x0d]),
1419                BranchNodeCompact::default(),
1420            );
1421            let encoded = bincode::serialize(&data).unwrap();
1422            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1423            assert_eq!(decoded, data);
1424
1425            data.trie_updates.storage_tries.insert(B256::default(), StorageTrieUpdates::default());
1426            let encoded = bincode::serialize(&data).unwrap();
1427            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1428            assert_eq!(decoded, data);
1429        }
1430
1431        #[test]
1432        fn test_storage_trie_updates_bincode_roundtrip() {
1433            #[serde_as]
1434            #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
1435            struct Data {
1436                #[serde_as(as = "serde_bincode_compat::updates::StorageTrieUpdates")]
1437                trie_updates: StorageTrieUpdates,
1438            }
1439
1440            let mut data = Data { trie_updates: StorageTrieUpdates::default() };
1441            let encoded = bincode::serialize(&data).unwrap();
1442            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1443            assert_eq!(decoded, data);
1444
1445            data.trie_updates
1446                .removed_nodes
1447                .insert(Nibbles::from_nibbles_unchecked([0x0b, 0x0e, 0x0e, 0x0f]));
1448            let encoded = bincode::serialize(&data).unwrap();
1449            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1450            assert_eq!(decoded, data);
1451
1452            data.trie_updates.storage_nodes.insert(
1453                Nibbles::from_nibbles_unchecked([0x0d, 0x0e, 0x0a, 0x0d]),
1454                BranchNodeCompact::default(),
1455            );
1456            let encoded = bincode::serialize(&data).unwrap();
1457            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1458            assert_eq!(decoded, data);
1459        }
1460
1461        #[test]
1462        fn test_trie_updates_sorted_bincode_roundtrip() {
1463            #[serde_as]
1464            #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
1465            struct Data {
1466                #[serde_as(as = "serde_bincode_compat::updates::TrieUpdatesSorted")]
1467                trie_updates: TrieUpdatesSorted,
1468            }
1469
1470            let mut data = Data { trie_updates: TrieUpdatesSorted::default() };
1471            let encoded = bincode::serialize(&data).unwrap();
1472            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1473            assert_eq!(decoded, data);
1474
1475            data.trie_updates.account_nodes.push((
1476                Nibbles::from_nibbles_unchecked([0x0d, 0x0e, 0x0a, 0x0d]),
1477                Some(BranchNodeCompact::default()),
1478            ));
1479            let encoded = bincode::serialize(&data).unwrap();
1480            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1481            assert_eq!(decoded, data);
1482
1483            data.trie_updates
1484                .account_nodes
1485                .push((Nibbles::from_nibbles_unchecked([0x0f, 0x0f, 0x0f, 0x0f]), None));
1486            let encoded = bincode::serialize(&data).unwrap();
1487            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1488            assert_eq!(decoded, data);
1489
1490            data.trie_updates
1491                .storage_tries
1492                .insert(B256::default(), StorageTrieUpdatesSorted::default());
1493            let encoded = bincode::serialize(&data).unwrap();
1494            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1495            assert_eq!(decoded, data);
1496        }
1497
1498        #[test]
1499        fn test_storage_trie_updates_sorted_bincode_roundtrip() {
1500            #[serde_as]
1501            #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
1502            struct Data {
1503                #[serde_as(as = "serde_bincode_compat::updates::StorageTrieUpdatesSorted")]
1504                trie_updates: StorageTrieUpdatesSorted,
1505            }
1506
1507            let mut data = Data { trie_updates: StorageTrieUpdatesSorted::default() };
1508            let encoded = bincode::serialize(&data).unwrap();
1509            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1510            assert_eq!(decoded, data);
1511
1512            data.trie_updates.storage_nodes.push((
1513                Nibbles::from_nibbles_unchecked([0x0d, 0x0e, 0x0a, 0x0d]),
1514                Some(BranchNodeCompact::default()),
1515            ));
1516            let encoded = bincode::serialize(&data).unwrap();
1517            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1518            assert_eq!(decoded, data);
1519
1520            data.trie_updates
1521                .storage_nodes
1522                .push((Nibbles::from_nibbles_unchecked([0x0a, 0x0a, 0x0a, 0x0a]), None));
1523            let encoded = bincode::serialize(&data).unwrap();
1524            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1525            assert_eq!(decoded, data);
1526
1527            data.trie_updates.is_deleted = true;
1528            let encoded = bincode::serialize(&data).unwrap();
1529            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1530            assert_eq!(decoded, data);
1531        }
1532    }
1533}
1534
1535#[cfg(all(test, feature = "serde"))]
1536mod serde_tests {
1537    use super::*;
1538
1539    #[test]
1540    fn test_trie_updates_serde_roundtrip() {
1541        let mut default_updates = TrieUpdates::default();
1542        let updates_serialized = serde_json::to_string(&default_updates).unwrap();
1543        let updates_deserialized: TrieUpdates = serde_json::from_str(&updates_serialized).unwrap();
1544        assert_eq!(updates_deserialized, default_updates);
1545
1546        default_updates
1547            .removed_nodes
1548            .insert(Nibbles::from_nibbles_unchecked([0x0b, 0x0e, 0x0e, 0x0f]));
1549        let updates_serialized = serde_json::to_string(&default_updates).unwrap();
1550        let updates_deserialized: TrieUpdates = serde_json::from_str(&updates_serialized).unwrap();
1551        assert_eq!(updates_deserialized, default_updates);
1552
1553        default_updates.account_nodes.insert(
1554            Nibbles::from_nibbles_unchecked([0x0d, 0x0e, 0x0a, 0x0d]),
1555            BranchNodeCompact::default(),
1556        );
1557        let updates_serialized = serde_json::to_string(&default_updates).unwrap();
1558        let updates_deserialized: TrieUpdates = serde_json::from_str(&updates_serialized).unwrap();
1559        assert_eq!(updates_deserialized, default_updates);
1560
1561        default_updates.storage_tries.insert(B256::default(), StorageTrieUpdates::default());
1562        let updates_serialized = serde_json::to_string(&default_updates).unwrap();
1563        let updates_deserialized: TrieUpdates = serde_json::from_str(&updates_serialized).unwrap();
1564        assert_eq!(updates_deserialized, default_updates);
1565    }
1566
1567    #[test]
1568    fn test_storage_trie_updates_serde_roundtrip() {
1569        let mut default_updates = StorageTrieUpdates::default();
1570        let updates_serialized = serde_json::to_string(&default_updates).unwrap();
1571        let updates_deserialized: StorageTrieUpdates =
1572            serde_json::from_str(&updates_serialized).unwrap();
1573        assert_eq!(updates_deserialized, default_updates);
1574
1575        default_updates
1576            .removed_nodes
1577            .insert(Nibbles::from_nibbles_unchecked([0x0b, 0x0e, 0x0e, 0x0f]));
1578        let updates_serialized = serde_json::to_string(&default_updates).unwrap();
1579        let updates_deserialized: StorageTrieUpdates =
1580            serde_json::from_str(&updates_serialized).unwrap();
1581        assert_eq!(updates_deserialized, default_updates);
1582
1583        default_updates.storage_nodes.insert(
1584            Nibbles::from_nibbles_unchecked([0x0d, 0x0e, 0x0a, 0x0d]),
1585            BranchNodeCompact::default(),
1586        );
1587        let updates_serialized = serde_json::to_string(&default_updates).unwrap();
1588        let updates_deserialized: StorageTrieUpdates =
1589            serde_json::from_str(&updates_serialized).unwrap();
1590        assert_eq!(updates_deserialized, default_updates);
1591    }
1592}