miden_objects/account/storage/
map.rs

1use alloc::collections::BTreeMap;
2
3use miden_crypto::merkle::EmptySubtreeRoots;
4use vm_core::EMPTY_WORD;
5
6use super::{
7    ByteReader, ByteWriter, Deserializable, DeserializationError, Digest, Serializable, Word,
8};
9use crate::{
10    Hasher,
11    account::StorageMapDelta,
12    crypto::{
13        hash::rpo::RpoDigest,
14        merkle::{InnerNodeInfo, LeafIndex, SMT_DEPTH, Smt, SmtLeaf, SmtProof},
15    },
16    errors::StorageMapError,
17};
18
19// ACCOUNT STORAGE MAP
20// ================================================================================================
21
22/// Empty storage map root.
23pub const EMPTY_STORAGE_MAP_ROOT: Digest = *EmptySubtreeRoots::entry(StorageMap::TREE_DEPTH, 0);
24
25/// An account storage map is a sparse merkle tree of depth [`Self::TREE_DEPTH`] (64).
26///
27/// It can be used to store a large amount of data in an account than would be otherwise possible
28/// using just the account's storage slots. This works by storing the root of the map's underlying
29/// SMT in one account storage slot. Each map entry is a leaf in the tree and its inclusion is
30/// proven while retrieving it (e.g. via `account::get_map_item`).
31///
32/// As a side-effect, this also means that _not all_ entries of the map have to be present at
33/// transaction execution time in order to access or modify the map. It is sufficient if _just_ the
34/// accessed/modified items are present in the advice provider.
35///
36/// Because the keys of the map are user-chosen and thus not necessarily uniformly distributed, the
37/// tree could be imbalanced and made less efficient. To mitigate that, the keys used in the
38/// storage map are hashed before they are inserted into the SMT, which creates a uniform
39/// distribution. The original keys are retained in a separate map. This causes redundancy but
40/// allows for introspection of the map, e.g. by querying the set of stored (original) keys which is
41/// useful in debugging and explorer scenarios.
42#[derive(Debug, Clone, PartialEq, Eq)]
43pub struct StorageMap {
44    /// The SMT where each key is the hashed original key.
45    smt: Smt,
46    /// The entries of the map where the key is the original user-chosen one.
47    ///
48    /// It is an invariant of this type that the map's entries are always consistent with the SMT's
49    /// entries and vice-versa.
50    map: BTreeMap<RpoDigest, Word>,
51}
52
53impl StorageMap {
54    // CONSTANTS
55    // --------------------------------------------------------------------------------------------
56
57    /// Depth of the storage tree.
58    pub const TREE_DEPTH: u8 = SMT_DEPTH;
59
60    /// The default value of empty leaves.
61    pub const EMPTY_VALUE: Word = Smt::EMPTY_VALUE;
62
63    // CONSTRUCTOR
64    // --------------------------------------------------------------------------------------------
65
66    /// Returns a new [StorageMap].
67    ///
68    /// All leaves in the returned tree are set to [Self::EMPTY_VALUE].
69    pub fn new() -> Self {
70        StorageMap { smt: Smt::new(), map: BTreeMap::new() }
71    }
72
73    /// Creates a new [`StorageMap`] from the provided key-value entries.
74    ///
75    /// # Errors
76    ///
77    /// Returns an error if:
78    /// - the provided entries contain multiple values for the same key.
79    pub fn with_entries<I: ExactSizeIterator<Item = (RpoDigest, Word)>>(
80        entries: impl IntoIterator<Item = (RpoDigest, Word), IntoIter = I>,
81    ) -> Result<Self, StorageMapError> {
82        let mut map = BTreeMap::new();
83
84        for (key, value) in entries {
85            if let Some(prev_value) = map.insert(key, value) {
86                return Err(StorageMapError::DuplicateKey {
87                    key,
88                    value0: prev_value,
89                    value1: value,
90                });
91            }
92        }
93
94        Ok(Self::from_btree_map(map))
95    }
96
97    /// Creates a new [`StorageMap`] from the given map. For internal use.
98    fn from_btree_map(map: BTreeMap<RpoDigest, Word>) -> Self {
99        let hashed_keys_iter = map.iter().map(|(key, value)| (Self::hash_key(*key), *value));
100        let smt = Smt::with_entries(hashed_keys_iter)
101            .expect("btree maps should not contain duplicate keys");
102
103        StorageMap { smt, map }
104    }
105
106    // PUBLIC ACCESSORS
107    // --------------------------------------------------------------------------------------------
108
109    /// Returns the root of the underlying sparse merkle tree.
110    pub fn root(&self) -> RpoDigest {
111        self.smt.root()
112    }
113
114    /// Returns the value corresponding to the key or [`Self::EMPTY_VALUE`] if the key is not
115    /// associated with a value.
116    pub fn get(&self, key: &RpoDigest) -> Word {
117        self.map.get(key).copied().unwrap_or_default()
118    }
119
120    /// Returns an opening of the leaf associated with `key`.
121    ///
122    /// Conceptually, an opening is a Merkle path to the leaf, as well as the leaf itself.
123    pub fn open(&self, key: &RpoDigest) -> SmtProof {
124        let key = Self::hash_key(*key);
125        self.smt.open(&key)
126    }
127
128    // ITERATORS
129    // --------------------------------------------------------------------------------------------
130
131    /// Returns an iterator over the leaves of the underlying [`Smt`].
132    pub fn leaves(&self) -> impl Iterator<Item = (LeafIndex<SMT_DEPTH>, &SmtLeaf)> {
133        self.smt.leaves() // Delegate to Smt's leaves method
134    }
135
136    /// Returns an iterator over the key value pairs of the map.
137    pub fn entries(&self) -> impl Iterator<Item = (&RpoDigest, &Word)> {
138        self.map.iter()
139    }
140
141    /// Returns an iterator over the inner nodes of the underlying [`Smt`].
142    pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
143        self.smt.inner_nodes() // Delegate to Smt's inner_nodes method
144    }
145
146    // DATA MUTATORS
147    // --------------------------------------------------------------------------------------------
148
149    /// Inserts or updates the given key value pair and returns the previous value, or
150    /// [`Self::EMPTY_VALUE`] if no entry was previously present.
151    ///
152    /// If the provided `value` is [`Self::EMPTY_VALUE`] the entry will be removed.
153    pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word {
154        if value == EMPTY_WORD {
155            self.map.remove(&key);
156        } else {
157            self.map.insert(key, value);
158        }
159
160        let key = Self::hash_key(key);
161        self.smt.insert(key, value) // Delegate to Smt's insert method
162    }
163
164    /// Applies the provided delta to this account storage.
165    pub fn apply_delta(&mut self, delta: &StorageMapDelta) -> Digest {
166        // apply the updated and cleared leaves to the storage map
167        for (&key, &value) in delta.leaves().iter() {
168            self.insert(key, value);
169        }
170
171        self.root()
172    }
173
174    /// Consumes the map and returns the underlying map of entries.
175    pub fn into_entries(self) -> BTreeMap<RpoDigest, Word> {
176        self.map
177    }
178
179    /// Hashes the given key to get the key of the SMT.
180    pub fn hash_key(key: RpoDigest) -> RpoDigest {
181        Hasher::hash_elements(key.as_elements())
182    }
183}
184
185impl Default for StorageMap {
186    fn default() -> Self {
187        Self::new()
188    }
189}
190
191// SERIALIZATION
192// ================================================================================================
193
194impl Serializable for StorageMap {
195    fn write_into<W: ByteWriter>(&self, target: &mut W) {
196        self.map.write_into(target);
197    }
198
199    fn get_size_hint(&self) -> usize {
200        self.smt.get_size_hint()
201    }
202}
203
204impl Deserializable for StorageMap {
205    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
206        let map = BTreeMap::read_from(source)?;
207        Ok(Self::from_btree_map(map))
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use assert_matches::assert_matches;
214    use miden_crypto::{Felt, hash::rpo::RpoDigest};
215
216    use super::{Deserializable, EMPTY_STORAGE_MAP_ROOT, Serializable, StorageMap, Word};
217    use crate::errors::StorageMapError;
218
219    #[test]
220    fn account_storage_serialization() {
221        // StorageMap for default types (empty map)
222        let storage_map_default = StorageMap::default();
223        let bytes = storage_map_default.to_bytes();
224        assert_eq!(storage_map_default, StorageMap::read_from_bytes(&bytes).unwrap());
225
226        // StorageMap with values
227        let storage_map_leaves_2: [(RpoDigest, Word); 2] = [
228            (
229                RpoDigest::new([Felt::new(101), Felt::new(102), Felt::new(103), Felt::new(104)]),
230                [Felt::new(1_u64), Felt::new(2_u64), Felt::new(3_u64), Felt::new(4_u64)],
231            ),
232            (
233                RpoDigest::new([Felt::new(105), Felt::new(106), Felt::new(107), Felt::new(108)]),
234                [Felt::new(5_u64), Felt::new(6_u64), Felt::new(7_u64), Felt::new(8_u64)],
235            ),
236        ];
237        let storage_map = StorageMap::with_entries(storage_map_leaves_2).unwrap();
238
239        let bytes = storage_map.to_bytes();
240        let deserialized_map = StorageMap::read_from_bytes(&bytes).unwrap();
241
242        assert_eq!(storage_map.root(), deserialized_map.root());
243
244        assert_eq!(storage_map, deserialized_map);
245    }
246
247    #[test]
248    fn test_empty_storage_map_constants() {
249        // If these values don't match, update the constants.
250        assert_eq!(StorageMap::default().root(), EMPTY_STORAGE_MAP_ROOT);
251    }
252
253    #[test]
254    fn account_storage_map_fails_on_duplicate_entries() {
255        // StorageMap with values
256        let storage_map_leaves_2: [(RpoDigest, Word); 2] = [
257            (
258                RpoDigest::new([Felt::new(101), Felt::new(102), Felt::new(103), Felt::new(104)]),
259                [Felt::new(1_u64), Felt::new(2_u64), Felt::new(3_u64), Felt::new(4_u64)],
260            ),
261            (
262                RpoDigest::new([Felt::new(101), Felt::new(102), Felt::new(103), Felt::new(104)]),
263                [Felt::new(5_u64), Felt::new(6_u64), Felt::new(7_u64), Felt::new(8_u64)],
264            ),
265        ];
266
267        let error = StorageMap::with_entries(storage_map_leaves_2).unwrap_err();
268        assert_matches!(error, StorageMapError::DuplicateKey { .. });
269    }
270}
miden_objects/account/storage/map.rs

miden_objects/account/storage/
map.rs
