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(|a, b| a.0.cmp(&b.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(|a, b| a.0.cmp(&b.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(|a, b| a.0.cmp(&b.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(|a, b| a.0.cmp(&b.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        const THRESHOLD: usize = 30;
636
637        let items: alloc::vec::Vec<_> = iter.into_iter().collect();
638        let k = items.len();
639
640        if k == 0 {
641            return Self::default().into();
642        }
643        if k == 1 {
644            return items.into_iter().next().expect("k == 1");
645        }
646
647        if k < THRESHOLD {
648            // Small k: extend loop, oldest-to-newest so newer overrides older.
649            let mut iter = items.iter().rev();
650            let mut acc = iter.next().expect("k > 0").as_ref().clone();
651            for next in iter {
652                acc.extend_ref_and_sort(next.as_ref());
653            }
654            return acc.into();
655        }
656
657        // Large k: k-way merge.
658        let account_nodes =
659            kway_merge_sorted(items.iter().map(|i| i.as_ref().account_nodes.as_slice()));
660
661        struct StorageAcc<'a> {
662            is_deleted: bool,
663            sealed: bool,
664            slices: Vec<&'a [(Nibbles, Option<BranchNodeCompact>)]>,
665        }
666
667        let mut acc: B256Map<StorageAcc<'_>> = B256Map::default();
668
669        for item in &items {
670            for (addr, storage) in &item.as_ref().storage_tries {
671                let entry = acc.entry(*addr).or_insert_with(|| StorageAcc {
672                    is_deleted: false,
673                    sealed: false,
674                    slices: Vec::new(),
675                });
676
677                if entry.sealed {
678                    continue;
679                }
680
681                entry.slices.push(storage.storage_nodes.as_slice());
682
683                if storage.is_deleted {
684                    entry.is_deleted = true;
685                    entry.sealed = true;
686                }
687            }
688        }
689
690        let storage_tries = acc
691            .into_iter()
692            .map(|(addr, entry)| {
693                let storage_nodes = kway_merge_sorted(entry.slices);
694                (addr, StorageTrieUpdatesSorted { is_deleted: entry.is_deleted, storage_nodes })
695            })
696            .collect();
697
698        Self { account_nodes, storage_tries }.into()
699    }
700}
701
702impl AsRef<Self> for TrieUpdatesSorted {
703    fn as_ref(&self) -> &Self {
704        self
705    }
706}
707
708impl From<TrieUpdatesSorted> for TrieUpdates {
709    fn from(sorted: TrieUpdatesSorted) -> Self {
710        let mut account_nodes = HashMap::default();
711        let mut removed_nodes = HashSet::default();
712
713        for (nibbles, node) in sorted.account_nodes {
714            if let Some(node) = node {
715                account_nodes.insert(nibbles, node);
716            } else {
717                removed_nodes.insert(nibbles);
718            }
719        }
720
721        let storage_tries = sorted
722            .storage_tries
723            .into_iter()
724            .map(|(address, storage)| (address, storage.into()))
725            .collect();
726
727        Self { account_nodes, removed_nodes, storage_tries }
728    }
729}
730
731/// Sorted storage trie updates reference used for serializing to file.
732#[derive(PartialEq, Eq, Clone, Default, Debug)]
733#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize))]
734pub struct StorageTrieUpdatesSortedRef<'a> {
735    /// Flag indicating whether the trie has been deleted/wiped.
736    pub is_deleted: bool,
737    /// Sorted collection of updated storage nodes with corresponding paths.
738    pub storage_nodes: BTreeMap<&'a Nibbles, &'a BranchNodeCompact>,
739    /// The set of removed storage node keys.
740    pub removed_nodes: BTreeSet<&'a Nibbles>,
741}
742
743/// Sorted trie updates used for lookups and insertions.
744#[derive(PartialEq, Eq, Clone, Default, Debug)]
745#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
746pub struct StorageTrieUpdatesSorted {
747    /// Flag indicating whether the trie has been deleted/wiped.
748    pub is_deleted: bool,
749    /// Sorted collection of updated storage nodes with corresponding paths. None indicates a node
750    /// is removed.
751    pub storage_nodes: Vec<(Nibbles, Option<BranchNodeCompact>)>,
752}
753
754impl StorageTrieUpdatesSorted {
755    /// Returns `true` if the trie was deleted.
756    pub const fn is_deleted(&self) -> bool {
757        self.is_deleted
758    }
759
760    /// Returns reference to updated storage nodes.
761    pub fn storage_nodes_ref(&self) -> &[(Nibbles, Option<BranchNodeCompact>)] {
762        &self.storage_nodes
763    }
764
765    /// Returns the total number of storage node updates.
766    pub const fn len(&self) -> usize {
767        self.storage_nodes.len()
768    }
769
770    /// Returns `true` if there are no storage node updates.
771    pub const fn is_empty(&self) -> bool {
772        self.storage_nodes.is_empty()
773    }
774
775    /// Extends the storage trie updates with another set of sorted updates.
776    ///
777    /// If `other` is marked as deleted, this will be marked as deleted and all nodes cleared.
778    /// Otherwise, nodes are merged with `other`'s values taking precedence for duplicates.
779    pub fn extend_ref(&mut self, other: &Self) {
780        if other.is_deleted {
781            self.is_deleted = true;
782            self.storage_nodes.clear();
783            self.storage_nodes.extend(other.storage_nodes.iter().cloned());
784            return;
785        }
786
787        // Extend storage nodes
788        extend_sorted_vec(&mut self.storage_nodes, &other.storage_nodes);
789        self.is_deleted = self.is_deleted || other.is_deleted;
790    }
791
792    /// Batch-merge sorted storage trie updates. Iterator yields **newest to oldest**.
793    /// If any update is deleted, older data is discarded.
794    pub fn merge_batch<'a>(updates: impl IntoIterator<Item = &'a Self>) -> Self {
795        let updates: Vec<_> = updates.into_iter().collect();
796        if updates.is_empty() {
797            return Self::default();
798        }
799
800        // Discard updates older than the first deletion since the trie was wiped at that point.
801        let del_idx = updates.iter().position(|u| u.is_deleted);
802        let relevant = del_idx.map_or(&updates[..], |idx| &updates[..=idx]);
803        let storage_nodes = kway_merge_sorted(relevant.iter().map(|u| u.storage_nodes.as_slice()));
804
805        Self { is_deleted: del_idx.is_some(), storage_nodes }
806    }
807}
808
809/// Excludes empty nibbles from the given iterator.
810fn exclude_empty(iter: impl IntoIterator<Item = Nibbles>) -> impl Iterator<Item = Nibbles> {
811    iter.into_iter().filter(|n| !n.is_empty())
812}
813
814/// Excludes empty nibbles from the given iterator of pairs where the nibbles are the key.
815fn exclude_empty_from_pair<V>(
816    iter: impl IntoIterator<Item = (Nibbles, V)>,
817) -> impl Iterator<Item = (Nibbles, V)> {
818    iter.into_iter().filter(|(n, _)| !n.is_empty())
819}
820
821impl From<StorageTrieUpdatesSorted> for StorageTrieUpdates {
822    fn from(sorted: StorageTrieUpdatesSorted) -> Self {
823        let mut storage_nodes = HashMap::default();
824        let mut removed_nodes = HashSet::default();
825
826        for (nibbles, node) in sorted.storage_nodes {
827            if let Some(node) = node {
828                storage_nodes.insert(nibbles, node);
829            } else {
830                removed_nodes.insert(nibbles);
831            }
832        }
833
834        Self { is_deleted: sorted.is_deleted, storage_nodes, removed_nodes }
835    }
836}
837
838#[cfg(test)]
839mod tests {
840    use super::*;
841    use alloy_primitives::B256;
842
843    #[test]
844    fn test_trie_updates_sorted_extend_ref() {
845        // Test extending with empty updates
846        let mut updates1 = TrieUpdatesSorted::default();
847        let updates2 = TrieUpdatesSorted::default();
848        updates1.extend_ref_and_sort(&updates2);
849        assert_eq!(updates1.account_nodes.len(), 0);
850        assert_eq!(updates1.storage_tries.len(), 0);
851
852        // Test extending account nodes
853        let mut updates1 = TrieUpdatesSorted {
854            account_nodes: vec![
855                (Nibbles::from_nibbles_unchecked([0x01]), Some(BranchNodeCompact::default())),
856                (Nibbles::from_nibbles_unchecked([0x03]), None),
857            ],
858            storage_tries: B256Map::default(),
859        };
860        let updates2 = TrieUpdatesSorted {
861            account_nodes: vec![
862                (Nibbles::from_nibbles_unchecked([0x02]), Some(BranchNodeCompact::default())),
863                (Nibbles::from_nibbles_unchecked([0x03]), Some(BranchNodeCompact::default())), /* Override */
864            ],
865            storage_tries: B256Map::default(),
866        };
867        updates1.extend_ref_and_sort(&updates2);
868        assert_eq!(updates1.account_nodes.len(), 3);
869        // Should be sorted: 0x01, 0x02, 0x03
870        assert_eq!(updates1.account_nodes[0].0, Nibbles::from_nibbles_unchecked([0x01]));
871        assert_eq!(updates1.account_nodes[1].0, Nibbles::from_nibbles_unchecked([0x02]));
872        assert_eq!(updates1.account_nodes[2].0, Nibbles::from_nibbles_unchecked([0x03]));
873        // 0x03 should have Some value from updates2 (override)
874        assert!(updates1.account_nodes[2].1.is_some());
875
876        // Test extending storage tries
877        let storage_trie1 = StorageTrieUpdatesSorted {
878            is_deleted: false,
879            storage_nodes: vec![(
880                Nibbles::from_nibbles_unchecked([0x0a]),
881                Some(BranchNodeCompact::default()),
882            )],
883        };
884        let storage_trie2 = StorageTrieUpdatesSorted {
885            is_deleted: false,
886            storage_nodes: vec![(Nibbles::from_nibbles_unchecked([0x0b]), None)],
887        };
888
889        let hashed_address1 = B256::from([1; 32]);
890        let hashed_address2 = B256::from([2; 32]);
891
892        let mut updates1 = TrieUpdatesSorted {
893            account_nodes: vec![],
894            storage_tries: B256Map::from_iter([(hashed_address1, storage_trie1.clone())]),
895        };
896        let updates2 = TrieUpdatesSorted {
897            account_nodes: vec![],
898            storage_tries: B256Map::from_iter([
899                (hashed_address1, storage_trie2),
900                (hashed_address2, storage_trie1),
901            ]),
902        };
903        updates1.extend_ref_and_sort(&updates2);
904        assert_eq!(updates1.storage_tries.len(), 2);
905        assert!(updates1.storage_tries.contains_key(&hashed_address1));
906        assert!(updates1.storage_tries.contains_key(&hashed_address2));
907        // Check that storage trie for hashed_address1 was extended
908        let merged_storage = &updates1.storage_tries[&hashed_address1];
909        assert_eq!(merged_storage.storage_nodes.len(), 2);
910    }
911
912    #[test]
913    fn test_storage_trie_updates_sorted_extend_ref_deleted() {
914        // Test case 1: Extending with a deleted storage trie that has nodes
915        let mut storage1 = StorageTrieUpdatesSorted {
916            is_deleted: false,
917            storage_nodes: vec![
918                (Nibbles::from_nibbles_unchecked([0x01]), Some(BranchNodeCompact::default())),
919                (Nibbles::from_nibbles_unchecked([0x02]), None),
920            ],
921        };
922
923        let storage2 = StorageTrieUpdatesSorted {
924            is_deleted: true,
925            storage_nodes: vec![
926                (Nibbles::from_nibbles_unchecked([0x03]), Some(BranchNodeCompact::default())),
927                (Nibbles::from_nibbles_unchecked([0x04]), None),
928            ],
929        };
930
931        storage1.extend_ref(&storage2);
932
933        // Should be marked as deleted
934        assert!(storage1.is_deleted);
935        // Original nodes should be cleared, but other's nodes should be added
936        assert_eq!(storage1.storage_nodes.len(), 2);
937        assert_eq!(storage1.storage_nodes[0].0, Nibbles::from_nibbles_unchecked([0x03]));
938        assert_eq!(storage1.storage_nodes[1].0, Nibbles::from_nibbles_unchecked([0x04]));
939
940        // Test case 2: Extending a deleted storage trie with more nodes
941        let mut storage3 = StorageTrieUpdatesSorted {
942            is_deleted: true,
943            storage_nodes: vec![(
944                Nibbles::from_nibbles_unchecked([0x05]),
945                Some(BranchNodeCompact::default()),
946            )],
947        };
948
949        let storage4 = StorageTrieUpdatesSorted {
950            is_deleted: true,
951            storage_nodes: vec![
952                (Nibbles::from_nibbles_unchecked([0x06]), Some(BranchNodeCompact::default())),
953                (Nibbles::from_nibbles_unchecked([0x07]), None),
954            ],
955        };
956
957        storage3.extend_ref(&storage4);
958
959        // Should remain deleted
960        assert!(storage3.is_deleted);
961        // Should have nodes from other (original cleared then extended)
962        assert_eq!(storage3.storage_nodes.len(), 2);
963        assert_eq!(storage3.storage_nodes[0].0, Nibbles::from_nibbles_unchecked([0x06]));
964        assert_eq!(storage3.storage_nodes[1].0, Nibbles::from_nibbles_unchecked([0x07]));
965    }
966
967    /// Test extending with storage tries adds both nodes and removed nodes correctly
968    #[test]
969    fn test_trie_updates_extend_from_sorted_with_storage_tries() {
970        let hashed_address = B256::from([1; 32]);
971
972        let mut updates = TrieUpdates::default();
973
974        let storage_trie = StorageTrieUpdatesSorted {
975            is_deleted: false,
976            storage_nodes: vec![
977                (Nibbles::from_nibbles_unchecked([0x0a]), Some(BranchNodeCompact::default())),
978                (Nibbles::from_nibbles_unchecked([0x0b]), None),
979            ],
980        };
981
982        let sorted = TrieUpdatesSorted {
983            account_nodes: vec![],
984            storage_tries: B256Map::from_iter([(hashed_address, storage_trie)]),
985        };
986
987        updates.extend_from_sorted(&sorted);
988
989        assert_eq!(updates.storage_tries.len(), 1);
990        let storage = updates.storage_tries.get(&hashed_address).unwrap();
991        assert!(!storage.is_deleted);
992        assert_eq!(storage.storage_nodes.len(), 1);
993        assert!(storage.removed_nodes.contains(&Nibbles::from_nibbles_unchecked([0x0b])));
994    }
995
996    /// Test deleted=true clears old storage nodes before adding new ones (critical edge case)
997    #[test]
998    fn test_trie_updates_extend_from_sorted_with_deleted_storage() {
999        let hashed_address = B256::from([1; 32]);
1000
1001        let mut updates = TrieUpdates::default();
1002        updates.storage_tries.insert(
1003            hashed_address,
1004            StorageTrieUpdates {
1005                is_deleted: false,
1006                storage_nodes: HashMap::from_iter([(
1007                    Nibbles::from_nibbles_unchecked([0x01]),
1008                    BranchNodeCompact::default(),
1009                )]),
1010                removed_nodes: Default::default(),
1011            },
1012        );
1013
1014        let storage_trie = StorageTrieUpdatesSorted {
1015            is_deleted: true,
1016            storage_nodes: vec![(
1017                Nibbles::from_nibbles_unchecked([0x0a]),
1018                Some(BranchNodeCompact::default()),
1019            )],
1020        };
1021
1022        let sorted = TrieUpdatesSorted {
1023            account_nodes: vec![],
1024            storage_tries: B256Map::from_iter([(hashed_address, storage_trie)]),
1025        };
1026
1027        updates.extend_from_sorted(&sorted);
1028
1029        let storage = updates.storage_tries.get(&hashed_address).unwrap();
1030        assert!(storage.is_deleted);
1031        // After deletion, old nodes should be cleared
1032        assert_eq!(storage.storage_nodes.len(), 1);
1033        assert!(storage.storage_nodes.contains_key(&Nibbles::from_nibbles_unchecked([0x0a])));
1034    }
1035
1036    /// Test non-deleted storage merges nodes and tracks removed nodes
1037    #[test]
1038    fn test_storage_trie_updates_extend_from_sorted_non_deleted() {
1039        let mut storage = StorageTrieUpdates {
1040            is_deleted: false,
1041            storage_nodes: HashMap::from_iter([(
1042                Nibbles::from_nibbles_unchecked([0x01]),
1043                BranchNodeCompact::default(),
1044            )]),
1045            removed_nodes: Default::default(),
1046        };
1047
1048        let sorted = StorageTrieUpdatesSorted {
1049            is_deleted: false,
1050            storage_nodes: vec![
1051                (Nibbles::from_nibbles_unchecked([0x02]), Some(BranchNodeCompact::default())),
1052                (Nibbles::from_nibbles_unchecked([0x03]), None),
1053            ],
1054        };
1055
1056        storage.extend_from_sorted(&sorted);
1057
1058        assert!(!storage.is_deleted);
1059        assert_eq!(storage.storage_nodes.len(), 2);
1060        assert!(storage.removed_nodes.contains(&Nibbles::from_nibbles_unchecked([0x03])));
1061    }
1062
1063    /// Test deleted=true clears old nodes before extending (edge case)
1064    #[test]
1065    fn test_storage_trie_updates_extend_from_sorted_deleted() {
1066        let mut storage = StorageTrieUpdates {
1067            is_deleted: false,
1068            storage_nodes: HashMap::from_iter([(
1069                Nibbles::from_nibbles_unchecked([0x01]),
1070                BranchNodeCompact::default(),
1071            )]),
1072            removed_nodes: Default::default(),
1073        };
1074
1075        let sorted = StorageTrieUpdatesSorted {
1076            is_deleted: true,
1077            storage_nodes: vec![(
1078                Nibbles::from_nibbles_unchecked([0x0a]),
1079                Some(BranchNodeCompact::default()),
1080            )],
1081        };
1082
1083        storage.extend_from_sorted(&sorted);
1084
1085        assert!(storage.is_deleted);
1086        // Old nodes should be cleared when deleted
1087        assert_eq!(storage.storage_nodes.len(), 1);
1088        assert!(storage.storage_nodes.contains_key(&Nibbles::from_nibbles_unchecked([0x0a])));
1089    }
1090
1091    /// Test empty nibbles are filtered out during conversion (edge case bug)
1092    #[test]
1093    fn test_trie_updates_extend_from_sorted_filters_empty_nibbles() {
1094        let mut updates = TrieUpdates::default();
1095
1096        let sorted = TrieUpdatesSorted {
1097            account_nodes: vec![
1098                (Nibbles::default(), Some(BranchNodeCompact::default())), // Empty nibbles
1099                (Nibbles::from_nibbles_unchecked([0x01]), Some(BranchNodeCompact::default())),
1100            ],
1101            storage_tries: B256Map::default(),
1102        };
1103
1104        updates.extend_from_sorted(&sorted);
1105
1106        // Empty nibbles should be filtered out
1107        assert_eq!(updates.account_nodes.len(), 1);
1108        assert!(updates.account_nodes.contains_key(&Nibbles::from_nibbles_unchecked([0x01])));
1109        assert!(!updates.account_nodes.contains_key(&Nibbles::default()));
1110    }
1111}
1112
1113/// Bincode-compatible trie updates type serde implementations.
1114#[cfg(feature = "serde-bincode-compat")]
1115pub mod serde_bincode_compat {
1116    use crate::{BranchNodeCompact, Nibbles};
1117    use alloc::borrow::Cow;
1118    use alloy_primitives::map::{B256Map, HashMap, HashSet};
1119    use serde::{Deserialize, Deserializer, Serialize, Serializer};
1120    use serde_with::{DeserializeAs, SerializeAs};
1121
1122    /// Bincode-compatible [`super::TrieUpdates`] serde implementation.
1123    ///
1124    /// Intended to use with the [`serde_with::serde_as`] macro in the following way:
1125    /// ```rust
1126    /// use reth_trie_common::{serde_bincode_compat, updates::TrieUpdates};
1127    /// use serde::{Deserialize, Serialize};
1128    /// use serde_with::serde_as;
1129    ///
1130    /// #[serde_as]
1131    /// #[derive(Serialize, Deserialize)]
1132    /// struct Data {
1133    ///     #[serde_as(as = "serde_bincode_compat::updates::TrieUpdates")]
1134    ///     trie_updates: TrieUpdates,
1135    /// }
1136    /// ```
1137    #[derive(Debug, Serialize, Deserialize)]
1138    pub struct TrieUpdates<'a> {
1139        account_nodes: Cow<'a, HashMap<Nibbles, BranchNodeCompact>>,
1140        removed_nodes: Cow<'a, HashSet<Nibbles>>,
1141        storage_tries: B256Map<StorageTrieUpdates<'a>>,
1142    }
1143
1144    impl<'a> From<&'a super::TrieUpdates> for TrieUpdates<'a> {
1145        fn from(value: &'a super::TrieUpdates) -> Self {
1146            Self {
1147                account_nodes: Cow::Borrowed(&value.account_nodes),
1148                removed_nodes: Cow::Borrowed(&value.removed_nodes),
1149                storage_tries: value.storage_tries.iter().map(|(k, v)| (*k, v.into())).collect(),
1150            }
1151        }
1152    }
1153
1154    impl<'a> From<TrieUpdates<'a>> for super::TrieUpdates {
1155        fn from(value: TrieUpdates<'a>) -> Self {
1156            Self {
1157                account_nodes: value.account_nodes.into_owned(),
1158                removed_nodes: value.removed_nodes.into_owned(),
1159                storage_tries: value
1160                    .storage_tries
1161                    .into_iter()
1162                    .map(|(k, v)| (k, v.into()))
1163                    .collect(),
1164            }
1165        }
1166    }
1167
1168    impl SerializeAs<super::TrieUpdates> for TrieUpdates<'_> {
1169        fn serialize_as<S>(source: &super::TrieUpdates, serializer: S) -> Result<S::Ok, S::Error>
1170        where
1171            S: Serializer,
1172        {
1173            TrieUpdates::from(source).serialize(serializer)
1174        }
1175    }
1176
1177    impl<'de> DeserializeAs<'de, super::TrieUpdates> for TrieUpdates<'de> {
1178        fn deserialize_as<D>(deserializer: D) -> Result<super::TrieUpdates, D::Error>
1179        where
1180            D: Deserializer<'de>,
1181        {
1182            TrieUpdates::deserialize(deserializer).map(Into::into)
1183        }
1184    }
1185
1186    /// Bincode-compatible [`super::StorageTrieUpdates`] serde implementation.
1187    ///
1188    /// Intended to use with the [`serde_with::serde_as`] macro in the following way:
1189    /// ```rust
1190    /// use reth_trie_common::{serde_bincode_compat, updates::StorageTrieUpdates};
1191    /// use serde::{Deserialize, Serialize};
1192    /// use serde_with::serde_as;
1193    ///
1194    /// #[serde_as]
1195    /// #[derive(Serialize, Deserialize)]
1196    /// struct Data {
1197    ///     #[serde_as(as = "serde_bincode_compat::updates::StorageTrieUpdates")]
1198    ///     trie_updates: StorageTrieUpdates,
1199    /// }
1200    /// ```
1201    #[derive(Debug, Serialize, Deserialize)]
1202    pub struct StorageTrieUpdates<'a> {
1203        is_deleted: bool,
1204        storage_nodes: Cow<'a, HashMap<Nibbles, BranchNodeCompact>>,
1205        removed_nodes: Cow<'a, HashSet<Nibbles>>,
1206    }
1207
1208    impl<'a> From<&'a super::StorageTrieUpdates> for StorageTrieUpdates<'a> {
1209        fn from(value: &'a super::StorageTrieUpdates) -> Self {
1210            Self {
1211                is_deleted: value.is_deleted,
1212                storage_nodes: Cow::Borrowed(&value.storage_nodes),
1213                removed_nodes: Cow::Borrowed(&value.removed_nodes),
1214            }
1215        }
1216    }
1217
1218    impl<'a> From<StorageTrieUpdates<'a>> for super::StorageTrieUpdates {
1219        fn from(value: StorageTrieUpdates<'a>) -> Self {
1220            Self {
1221                is_deleted: value.is_deleted,
1222                storage_nodes: value.storage_nodes.into_owned(),
1223                removed_nodes: value.removed_nodes.into_owned(),
1224            }
1225        }
1226    }
1227
1228    impl SerializeAs<super::StorageTrieUpdates> for StorageTrieUpdates<'_> {
1229        fn serialize_as<S>(
1230            source: &super::StorageTrieUpdates,
1231            serializer: S,
1232        ) -> Result<S::Ok, S::Error>
1233        where
1234            S: Serializer,
1235        {
1236            StorageTrieUpdates::from(source).serialize(serializer)
1237        }
1238    }
1239
1240    impl<'de> DeserializeAs<'de, super::StorageTrieUpdates> for StorageTrieUpdates<'de> {
1241        fn deserialize_as<D>(deserializer: D) -> Result<super::StorageTrieUpdates, D::Error>
1242        where
1243            D: Deserializer<'de>,
1244        {
1245            StorageTrieUpdates::deserialize(deserializer).map(Into::into)
1246        }
1247    }
1248
1249    /// Bincode-compatible [`super::TrieUpdatesSorted`] serde implementation.
1250    ///
1251    /// Intended to use with the [`serde_with::serde_as`] macro in the following way:
1252    /// ```rust
1253    /// use reth_trie_common::{serde_bincode_compat, updates::TrieUpdatesSorted};
1254    /// use serde::{Deserialize, Serialize};
1255    /// use serde_with::serde_as;
1256    ///
1257    /// #[serde_as]
1258    /// #[derive(Serialize, Deserialize)]
1259    /// struct Data {
1260    ///     #[serde_as(as = "serde_bincode_compat::updates::TrieUpdatesSorted")]
1261    ///     trie_updates: TrieUpdatesSorted,
1262    /// }
1263    /// ```
1264    #[derive(Debug, Serialize, Deserialize)]
1265    pub struct TrieUpdatesSorted<'a> {
1266        account_nodes: Cow<'a, [(Nibbles, Option<BranchNodeCompact>)]>,
1267        storage_tries: B256Map<StorageTrieUpdatesSorted<'a>>,
1268    }
1269
1270    impl<'a> From<&'a super::TrieUpdatesSorted> for TrieUpdatesSorted<'a> {
1271        fn from(value: &'a super::TrieUpdatesSorted) -> Self {
1272            Self {
1273                account_nodes: Cow::Borrowed(&value.account_nodes),
1274                storage_tries: value.storage_tries.iter().map(|(k, v)| (*k, v.into())).collect(),
1275            }
1276        }
1277    }
1278
1279    impl<'a> From<TrieUpdatesSorted<'a>> for super::TrieUpdatesSorted {
1280        fn from(value: TrieUpdatesSorted<'a>) -> Self {
1281            Self {
1282                account_nodes: value.account_nodes.into_owned(),
1283                storage_tries: value
1284                    .storage_tries
1285                    .into_iter()
1286                    .map(|(k, v)| (k, v.into()))
1287                    .collect(),
1288            }
1289        }
1290    }
1291
1292    impl SerializeAs<super::TrieUpdatesSorted> for TrieUpdatesSorted<'_> {
1293        fn serialize_as<S>(
1294            source: &super::TrieUpdatesSorted,
1295            serializer: S,
1296        ) -> Result<S::Ok, S::Error>
1297        where
1298            S: Serializer,
1299        {
1300            TrieUpdatesSorted::from(source).serialize(serializer)
1301        }
1302    }
1303
1304    impl<'de> DeserializeAs<'de, super::TrieUpdatesSorted> for TrieUpdatesSorted<'de> {
1305        fn deserialize_as<D>(deserializer: D) -> Result<super::TrieUpdatesSorted, D::Error>
1306        where
1307            D: Deserializer<'de>,
1308        {
1309            TrieUpdatesSorted::deserialize(deserializer).map(Into::into)
1310        }
1311    }
1312
1313    /// Bincode-compatible [`super::StorageTrieUpdatesSorted`] serde implementation.
1314    ///
1315    /// Intended to use with the [`serde_with::serde_as`] macro in the following way:
1316    /// ```rust
1317    /// use reth_trie_common::{serde_bincode_compat, updates::StorageTrieUpdatesSorted};
1318    /// use serde::{Deserialize, Serialize};
1319    /// use serde_with::serde_as;
1320    ///
1321    /// #[serde_as]
1322    /// #[derive(Serialize, Deserialize)]
1323    /// struct Data {
1324    ///     #[serde_as(as = "serde_bincode_compat::updates::StorageTrieUpdatesSorted")]
1325    ///     trie_updates: StorageTrieUpdatesSorted,
1326    /// }
1327    /// ```
1328    #[derive(Debug, Serialize, Deserialize)]
1329    pub struct StorageTrieUpdatesSorted<'a> {
1330        is_deleted: bool,
1331        storage_nodes: Cow<'a, [(Nibbles, Option<BranchNodeCompact>)]>,
1332    }
1333
1334    impl<'a> From<&'a super::StorageTrieUpdatesSorted> for StorageTrieUpdatesSorted<'a> {
1335        fn from(value: &'a super::StorageTrieUpdatesSorted) -> Self {
1336            Self {
1337                is_deleted: value.is_deleted,
1338                storage_nodes: Cow::Borrowed(&value.storage_nodes),
1339            }
1340        }
1341    }
1342
1343    impl<'a> From<StorageTrieUpdatesSorted<'a>> for super::StorageTrieUpdatesSorted {
1344        fn from(value: StorageTrieUpdatesSorted<'a>) -> Self {
1345            Self { is_deleted: value.is_deleted, storage_nodes: value.storage_nodes.into_owned() }
1346        }
1347    }
1348
1349    impl SerializeAs<super::StorageTrieUpdatesSorted> for StorageTrieUpdatesSorted<'_> {
1350        fn serialize_as<S>(
1351            source: &super::StorageTrieUpdatesSorted,
1352            serializer: S,
1353        ) -> Result<S::Ok, S::Error>
1354        where
1355            S: Serializer,
1356        {
1357            StorageTrieUpdatesSorted::from(source).serialize(serializer)
1358        }
1359    }
1360
1361    impl<'de> DeserializeAs<'de, super::StorageTrieUpdatesSorted> for StorageTrieUpdatesSorted<'de> {
1362        fn deserialize_as<D>(deserializer: D) -> Result<super::StorageTrieUpdatesSorted, D::Error>
1363        where
1364            D: Deserializer<'de>,
1365        {
1366            StorageTrieUpdatesSorted::deserialize(deserializer).map(Into::into)
1367        }
1368    }
1369
1370    #[cfg(test)]
1371    mod tests {
1372        use crate::{
1373            serde_bincode_compat,
1374            updates::{
1375                StorageTrieUpdates, StorageTrieUpdatesSorted, TrieUpdates, TrieUpdatesSorted,
1376            },
1377            BranchNodeCompact, Nibbles,
1378        };
1379        use alloy_primitives::B256;
1380        use serde::{Deserialize, Serialize};
1381        use serde_with::serde_as;
1382
1383        #[test]
1384        fn test_trie_updates_bincode_roundtrip() {
1385            #[serde_as]
1386            #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
1387            struct Data {
1388                #[serde_as(as = "serde_bincode_compat::updates::TrieUpdates")]
1389                trie_updates: TrieUpdates,
1390            }
1391
1392            let mut data = Data { trie_updates: TrieUpdates::default() };
1393            let encoded = bincode::serialize(&data).unwrap();
1394            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1395            assert_eq!(decoded, data);
1396
1397            data.trie_updates
1398                .removed_nodes
1399                .insert(Nibbles::from_nibbles_unchecked([0x0b, 0x0e, 0x0e, 0x0f]));
1400            let encoded = bincode::serialize(&data).unwrap();
1401            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1402            assert_eq!(decoded, data);
1403
1404            data.trie_updates.account_nodes.insert(
1405                Nibbles::from_nibbles_unchecked([0x0d, 0x0e, 0x0a, 0x0d]),
1406                BranchNodeCompact::default(),
1407            );
1408            let encoded = bincode::serialize(&data).unwrap();
1409            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1410            assert_eq!(decoded, data);
1411
1412            data.trie_updates.storage_tries.insert(B256::default(), StorageTrieUpdates::default());
1413            let encoded = bincode::serialize(&data).unwrap();
1414            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1415            assert_eq!(decoded, data);
1416        }
1417
1418        #[test]
1419        fn test_storage_trie_updates_bincode_roundtrip() {
1420            #[serde_as]
1421            #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
1422            struct Data {
1423                #[serde_as(as = "serde_bincode_compat::updates::StorageTrieUpdates")]
1424                trie_updates: StorageTrieUpdates,
1425            }
1426
1427            let mut data = Data { trie_updates: StorageTrieUpdates::default() };
1428            let encoded = bincode::serialize(&data).unwrap();
1429            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1430            assert_eq!(decoded, data);
1431
1432            data.trie_updates
1433                .removed_nodes
1434                .insert(Nibbles::from_nibbles_unchecked([0x0b, 0x0e, 0x0e, 0x0f]));
1435            let encoded = bincode::serialize(&data).unwrap();
1436            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1437            assert_eq!(decoded, data);
1438
1439            data.trie_updates.storage_nodes.insert(
1440                Nibbles::from_nibbles_unchecked([0x0d, 0x0e, 0x0a, 0x0d]),
1441                BranchNodeCompact::default(),
1442            );
1443            let encoded = bincode::serialize(&data).unwrap();
1444            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1445            assert_eq!(decoded, data);
1446        }
1447
1448        #[test]
1449        fn test_trie_updates_sorted_bincode_roundtrip() {
1450            #[serde_as]
1451            #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
1452            struct Data {
1453                #[serde_as(as = "serde_bincode_compat::updates::TrieUpdatesSorted")]
1454                trie_updates: TrieUpdatesSorted,
1455            }
1456
1457            let mut data = Data { trie_updates: TrieUpdatesSorted::default() };
1458            let encoded = bincode::serialize(&data).unwrap();
1459            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1460            assert_eq!(decoded, data);
1461
1462            data.trie_updates.account_nodes.push((
1463                Nibbles::from_nibbles_unchecked([0x0d, 0x0e, 0x0a, 0x0d]),
1464                Some(BranchNodeCompact::default()),
1465            ));
1466            let encoded = bincode::serialize(&data).unwrap();
1467            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1468            assert_eq!(decoded, data);
1469
1470            data.trie_updates
1471                .account_nodes
1472                .push((Nibbles::from_nibbles_unchecked([0x0f, 0x0f, 0x0f, 0x0f]), None));
1473            let encoded = bincode::serialize(&data).unwrap();
1474            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1475            assert_eq!(decoded, data);
1476
1477            data.trie_updates
1478                .storage_tries
1479                .insert(B256::default(), StorageTrieUpdatesSorted::default());
1480            let encoded = bincode::serialize(&data).unwrap();
1481            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1482            assert_eq!(decoded, data);
1483        }
1484
1485        #[test]
1486        fn test_storage_trie_updates_sorted_bincode_roundtrip() {
1487            #[serde_as]
1488            #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
1489            struct Data {
1490                #[serde_as(as = "serde_bincode_compat::updates::StorageTrieUpdatesSorted")]
1491                trie_updates: StorageTrieUpdatesSorted,
1492            }
1493
1494            let mut data = Data { trie_updates: StorageTrieUpdatesSorted::default() };
1495            let encoded = bincode::serialize(&data).unwrap();
1496            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1497            assert_eq!(decoded, data);
1498
1499            data.trie_updates.storage_nodes.push((
1500                Nibbles::from_nibbles_unchecked([0x0d, 0x0e, 0x0a, 0x0d]),
1501                Some(BranchNodeCompact::default()),
1502            ));
1503            let encoded = bincode::serialize(&data).unwrap();
1504            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1505            assert_eq!(decoded, data);
1506
1507            data.trie_updates
1508                .storage_nodes
1509                .push((Nibbles::from_nibbles_unchecked([0x0a, 0x0a, 0x0a, 0x0a]), None));
1510            let encoded = bincode::serialize(&data).unwrap();
1511            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1512            assert_eq!(decoded, data);
1513
1514            data.trie_updates.is_deleted = true;
1515            let encoded = bincode::serialize(&data).unwrap();
1516            let decoded: Data = bincode::deserialize(&encoded).unwrap();
1517            assert_eq!(decoded, data);
1518        }
1519    }
1520}
1521
1522#[cfg(all(test, feature = "serde"))]
1523mod serde_tests {
1524    use super::*;
1525
1526    #[test]
1527    fn test_trie_updates_serde_roundtrip() {
1528        let mut default_updates = TrieUpdates::default();
1529        let updates_serialized = serde_json::to_string(&default_updates).unwrap();
1530        let updates_deserialized: TrieUpdates = serde_json::from_str(&updates_serialized).unwrap();
1531        assert_eq!(updates_deserialized, default_updates);
1532
1533        default_updates
1534            .removed_nodes
1535            .insert(Nibbles::from_nibbles_unchecked([0x0b, 0x0e, 0x0e, 0x0f]));
1536        let updates_serialized = serde_json::to_string(&default_updates).unwrap();
1537        let updates_deserialized: TrieUpdates = serde_json::from_str(&updates_serialized).unwrap();
1538        assert_eq!(updates_deserialized, default_updates);
1539
1540        default_updates.account_nodes.insert(
1541            Nibbles::from_nibbles_unchecked([0x0d, 0x0e, 0x0a, 0x0d]),
1542            BranchNodeCompact::default(),
1543        );
1544        let updates_serialized = serde_json::to_string(&default_updates).unwrap();
1545        let updates_deserialized: TrieUpdates = serde_json::from_str(&updates_serialized).unwrap();
1546        assert_eq!(updates_deserialized, default_updates);
1547
1548        default_updates.storage_tries.insert(B256::default(), StorageTrieUpdates::default());
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
1554    #[test]
1555    fn test_storage_trie_updates_serde_roundtrip() {
1556        let mut default_updates = StorageTrieUpdates::default();
1557        let updates_serialized = serde_json::to_string(&default_updates).unwrap();
1558        let updates_deserialized: StorageTrieUpdates =
1559            serde_json::from_str(&updates_serialized).unwrap();
1560        assert_eq!(updates_deserialized, default_updates);
1561
1562        default_updates
1563            .removed_nodes
1564            .insert(Nibbles::from_nibbles_unchecked([0x0b, 0x0e, 0x0e, 0x0f]));
1565        let updates_serialized = serde_json::to_string(&default_updates).unwrap();
1566        let updates_deserialized: StorageTrieUpdates =
1567            serde_json::from_str(&updates_serialized).unwrap();
1568        assert_eq!(updates_deserialized, default_updates);
1569
1570        default_updates.storage_nodes.insert(
1571            Nibbles::from_nibbles_unchecked([0x0d, 0x0e, 0x0a, 0x0d]),
1572            BranchNodeCompact::default(),
1573        );
1574        let updates_serialized = serde_json::to_string(&default_updates).unwrap();
1575        let updates_deserialized: StorageTrieUpdates =
1576            serde_json::from_str(&updates_serialized).unwrap();
1577        assert_eq!(updates_deserialized, default_updates);
1578    }
1579}