diff --git a/crates/trie/trie/src/updates.rs b/crates/trie/trie/src/updates.rs index 9cb42fd68c..365211f2fe 100644 --- a/crates/trie/trie/src/updates.rs +++ b/crates/trie/trie/src/updates.rs @@ -1,5 +1,7 @@ use crate::{walker::TrieWalker, BranchNodeCompact, HashBuilder, Nibbles}; use alloy_primitives::B256; +#[cfg(feature = "serde")] +use serde::{ser::SerializeMap, Serialize, Serializer}; use std::collections::{HashMap, HashSet}; /// The aggregation of trie updates. @@ -113,6 +115,7 @@ pub struct StorageTrieUpdates { /// Flag indicating whether the trie was deleted. pub(crate) is_deleted: bool, /// Collection of updated storage trie nodes. + #[cfg_attr(feature = "serde", serde(serialize_with = "serialize_nibbles_map"))] pub(crate) storage_nodes: HashMap, /// Collection of removed storage trie nodes. pub(crate) removed_nodes: HashSet, @@ -216,6 +219,27 @@ impl StorageTrieUpdates { } } +/// Serializes any [`HashMap`] that uses [`Nibbles`] as keys, by using the hex-encoded packed +/// representation. +/// +/// This also sorts the map's keys before encoding and serializing. +#[cfg(feature = "serde")] +fn serialize_nibbles_map(map: &HashMap, serializer: S) -> Result +where + S: Serializer, + T: Serialize, +{ + let mut map_serializer = serializer.serialize_map(Some(map.len()))?; + let mut storage_nodes = Vec::from_iter(map); + storage_nodes.sort_unstable_by(|a, b| a.0.cmp(b.0)); + for (k, v) in storage_nodes { + // pack, then hex encode the Nibbles + let packed = reth_primitives::hex::encode(k.pack()); + map_serializer.serialize_entry(&packed, &v)?; + } + map_serializer.end() +} + /// Sorted trie updates used for lookups and insertions. #[derive(PartialEq, Eq, Clone, Default, Debug)] pub struct TrieUpdatesSorted { @@ -277,3 +301,22 @@ fn exclude_empty_from_pair( ) -> impl Iterator { iter.into_iter().filter(|(n, _)| !n.is_empty()) } + +#[cfg(all(test, feature = "serde"))] +mod tests { + use super::*; + + #[test] + fn test_serialize_works() { + let mut default_updates = StorageTrieUpdates::default(); + let _updates_string = serde_json::to_string(&default_updates).unwrap(); + + default_updates.removed_nodes.insert(Nibbles::from_vec(vec![0x0b, 0x0e, 0x0e, 0x0f])); + let _updates_string = serde_json::to_string(&default_updates).unwrap(); + + default_updates + .storage_nodes + .insert(Nibbles::from_vec(vec![0x0b, 0x0e, 0x0f]), BranchNodeCompact::default()); + let _updates_string = serde_json::to_string(&default_updates).unwrap(); + } +}