diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 5dbe289b..f59ea6ca 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -2,7 +2,11 @@ use super::{ EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD, }; -use alloc::collections::{BTreeMap, BTreeSet}; +use alloc::{ + collections::{BTreeMap, BTreeSet}, + string::ToString, + vec::Vec, +}; mod error; pub use error::{SmtLeafError, SmtProofError}; @@ -12,6 +16,7 @@ pub use leaf::SmtLeaf; mod proof; pub use proof::SmtProof; +use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; #[cfg(test)] mod tests; @@ -294,3 +299,62 @@ impl From<&RpoDigest> for LeafIndex { Word::from(value).into() } } + +// SERIALIZATION +// ================================================================================================ + +impl Serializable for Smt { + fn write_into(&self, target: &mut W) { + let filled_leaves = self.entries().collect::>(); + + // Write the number of filled leaves for this map + target.write_u8(filled_leaves.len() as u8); + + // Write each (key, value) pair + for (key, value) in filled_leaves { + target.write(key); + target.write(value); + } + } +} + +impl Deserializable for Smt { + fn read_from(source: &mut R) -> Result { + // read the number of filled leaves for this StorageMap + let num_filled_leaves = source.read_u8()?; + let mut entries = Vec::with_capacity(num_filled_leaves as usize); + + for _ in 0..num_filled_leaves { + let key = source.read()?; + let value = source.read()?; + entries.push((key, value)); + } + + Self::with_entries(entries) + .map_err(|err| DeserializationError::InvalidValue(err.to_string())) + } +} + +#[test] +fn test_smt_serialization_deserialization() { + // SMT for default types (empty map) + let storage_map_default = Smt::default(); + let bytes = storage_map_default.to_bytes(); + assert_eq!(storage_map_default, Smt::read_from_bytes(&bytes).unwrap()); + + // SMT with values + let storage_map_leaves_2: [(RpoDigest, Word); 2] = [ + ( + RpoDigest::new([Felt::new(101), Felt::new(102), Felt::new(103), Felt::new(104)]), + [Felt::new(1_u64), Felt::new(2_u64), Felt::new(3_u64), Felt::new(4_u64)], + ), + ( + RpoDigest::new([Felt::new(105), Felt::new(106), Felt::new(107), Felt::new(108)]), + [Felt::new(5_u64), Felt::new(6_u64), Felt::new(7_u64), Felt::new(8_u64)], + ), + ]; + let storage_map = Smt::with_entries(storage_map_leaves_2).unwrap(); + + let bytes = storage_map.to_bytes(); + assert_eq!(storage_map, Smt::read_from_bytes(&bytes).unwrap()); +}