Skip to main content

miden_protocol/account/patch/storage/
map_patch.rs

1use alloc::collections::BTreeMap;
2
3use super::slot_patch::MergeOutcome;
4use crate::Word;
5use crate::account::{StorageMap, StorageMapKey, StoragePatchOperation, StorageSlotName};
6use crate::errors::AccountPatchError;
7use crate::utils::serde::{
8    ByteReader,
9    ByteWriter,
10    Deserializable,
11    DeserializationError,
12    Serializable,
13};
14
15// STORAGE MAP PATCH
16// ================================================================================================
17
18/// The patch of a single map storage slot.
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub enum StorageMapPatch {
21    /// Records the creation of the map with the given entries.
22    ///
23    /// An empty entry set is meaningful to represent the creation of an empty map.
24    Create { entries: StorageMapPatchEntries },
25
26    /// Records that the entries of an existing map were changed.
27    ///
28    /// The entries should be non-empty, since empty entries are a no-op.
29    Update { entries: StorageMapPatchEntries },
30
31    /// Records that the map slot was removed.
32    Remove,
33}
34
35impl StorageMapPatch {
36    // CONSTANTS
37    // ----------------------------------------------------------------------------------------
38
39    const CREATE: u8 = 0;
40    const UPDATE: u8 = 1;
41    const REMOVE: u8 = 2;
42
43    // ACCESSORS
44    // ----------------------------------------------------------------------------------------
45
46    /// Returns the entries of the map patch for [`StorageMapPatch::Create`] and
47    /// [`StorageMapPatch::Update`], or `None` for [`StorageMapPatch::Remove`].
48    pub fn entries(&self) -> Option<&StorageMapPatchEntries> {
49        match self {
50            StorageMapPatch::Create { entries } | StorageMapPatch::Update { entries } => {
51                Some(entries)
52            },
53            StorageMapPatch::Remove => None,
54        }
55    }
56
57    /// Consumes self and returns the entries of the map patch for [`StorageMapPatch::Create`] and
58    /// [`StorageMapPatch::Update`], or `None` for [`StorageMapPatch::Remove`].
59    pub fn into_entries(self) -> Option<StorageMapPatchEntries> {
60        match self {
61            StorageMapPatch::Create { entries } | StorageMapPatch::Update { entries } => {
62                Some(entries)
63            },
64            StorageMapPatch::Remove => None,
65        }
66    }
67
68    /// Returns the [`StoragePatchOperation`] that this patch represents.
69    pub fn patch_op(&self) -> StoragePatchOperation {
70        match self {
71            StorageMapPatch::Create { .. } => StoragePatchOperation::Create,
72            StorageMapPatch::Update { .. } => StoragePatchOperation::Update,
73            StorageMapPatch::Remove => StoragePatchOperation::Remove,
74        }
75    }
76
77    // HELPERS
78    // ----------------------------------------------------------------------------------------
79
80    /// Merges `other` into `self`, with `other` taking precedence.
81    ///
82    /// A map that was created and then updated remains created (with the merged entries). A map
83    /// that was created and then removed cancels out, signalled via [`MergeOutcome::Remove`].
84    pub(super) fn merge(
85        &mut self,
86        slot_name: &StorageSlotName,
87        other: Self,
88    ) -> Result<MergeOutcome, AccountPatchError> {
89        match (self, other) {
90            // (Create, _) patterns
91            // ------------------------------------------------------------------------------------
92            (StorageMapPatch::Create { .. }, StorageMapPatch::Create { .. }) => {
93                return Err(AccountPatchError::StoragePatchMergeDoubleCreate(slot_name.clone()));
94            },
95            (
96                StorageMapPatch::Create { entries: current },
97                StorageMapPatch::Update { entries: incoming },
98            ) => current.merge(incoming),
99            (StorageMapPatch::Create { .. }, StorageMapPatch::Remove) => {
100                return Ok(MergeOutcome::Remove);
101            },
102
103            // (Update, _) patterns
104            // ------------------------------------------------------------------------------------
105            (StorageMapPatch::Update { .. }, StorageMapPatch::Create { .. }) => {
106                return Err(AccountPatchError::StoragePatchMergeCreateAfterUpdate(
107                    slot_name.clone(),
108                ));
109            },
110            (
111                StorageMapPatch::Update { entries: current },
112                StorageMapPatch::Update { entries: incoming },
113            ) => current.merge(incoming),
114            (current @ StorageMapPatch::Update { .. }, StorageMapPatch::Remove) => {
115                *current = StorageMapPatch::Remove
116            },
117
118            // (Remove, _) patterns
119            // ------------------------------------------------------------------------------------
120            (current @ StorageMapPatch::Remove, incoming @ StorageMapPatch::Create { .. }) => {
121                *current = incoming;
122            },
123            (StorageMapPatch::Remove, StorageMapPatch::Update { .. }) => {
124                return Err(AccountPatchError::StoragePatchMergeUpdateAfterRemove(
125                    slot_name.clone(),
126                ));
127            },
128            (StorageMapPatch::Remove, StorageMapPatch::Remove) => {
129                return Err(AccountPatchError::StoragePatchMergeDoubleRemove(slot_name.clone()));
130            },
131        }
132
133        Ok(MergeOutcome::Keep)
134    }
135}
136
137impl Serializable for StorageMapPatch {
138    fn write_into<W: ByteWriter>(&self, target: &mut W) {
139        match self {
140            StorageMapPatch::Create { entries } => {
141                target.write_u8(Self::CREATE);
142                target.write(entries);
143            },
144            StorageMapPatch::Update { entries } => {
145                target.write_u8(Self::UPDATE);
146                target.write(entries);
147            },
148            StorageMapPatch::Remove => {
149                target.write_u8(Self::REMOVE);
150            },
151        }
152    }
153
154    fn get_size_hint(&self) -> usize {
155        let tag_size = 0u8.get_size_hint();
156        match self {
157            StorageMapPatch::Create { entries } | StorageMapPatch::Update { entries } => {
158                tag_size + entries.get_size_hint()
159            },
160            StorageMapPatch::Remove => tag_size,
161        }
162    }
163}
164
165impl Deserializable for StorageMapPatch {
166    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
167        match source.read_u8()? {
168            Self::CREATE => Ok(Self::Create { entries: source.read()? }),
169            Self::UPDATE => Ok(Self::Update { entries: source.read()? }),
170            Self::REMOVE => Ok(Self::Remove),
171            other => Err(DeserializationError::InvalidValue(format!(
172                "unknown storage map patch variant {other}"
173            ))),
174        }
175    }
176}
177
178// STORAGE MAP PATCH ENTRIES
179// ================================================================================================
180
181/// The changed entries of a storage map, represented as a map of changed item key
182/// ([`StorageMapKey`]) to value ([`Word`]). For cleared items the value is [`Word::empty`].
183#[derive(Clone, Debug, Default, PartialEq, Eq)]
184pub struct StorageMapPatchEntries(BTreeMap<StorageMapKey, Word>);
185
186impl StorageMapPatchEntries {
187    /// Creates a new, empty set of map patch entries.
188    pub fn new() -> Self {
189        Self(BTreeMap::new())
190    }
191
192    /// Creates a new set of map patch entries from the provided map.
193    pub fn from_raw(entries: BTreeMap<StorageMapKey, Word>) -> Self {
194        Self(entries)
195    }
196
197    /// Returns the number of changed entries.
198    pub fn num_entries(&self) -> usize {
199        self.0.len()
200    }
201
202    /// Returns a reference to the changed entries.
203    ///
204    /// Note that the returned key is the [`StorageMapKey`].
205    pub fn as_map(&self) -> &BTreeMap<StorageMapKey, Word> {
206        &self.0
207    }
208
209    /// Inserts an entry.
210    pub fn insert(&mut self, key: StorageMapKey, value: Word) {
211        self.0.insert(key, value);
212    }
213
214    /// Returns true if there are no changed entries.
215    pub fn is_empty(&self) -> bool {
216        self.0.is_empty()
217    }
218
219    /// Returns an iterator over the keys of the cleared entries, i.e. those whose value is
220    /// [`Word::empty`].
221    fn cleared_entries(&self) -> impl Iterator<Item = &StorageMapKey> + Clone {
222        self.0.iter().filter(|(_, value)| value.is_empty()).map(|(key, _)| key)
223    }
224
225    /// Returns an iterator over the updated entries, i.e. those whose value is not [`Word::empty`].
226    fn updated_entries(&self) -> impl Iterator<Item = (&StorageMapKey, &Word)> + Clone {
227        self.0.iter().filter(|(_, value)| !value.is_empty())
228    }
229
230    /// Merges `other` into these entries, with the entries of `other` taking precedence.
231    fn merge(&mut self, other: Self) {
232        self.0.extend(other.0);
233    }
234
235    /// Returns a mutable reference to the underlying map.
236    pub fn as_map_mut(&mut self) -> &mut BTreeMap<StorageMapKey, Word> {
237        &mut self.0
238    }
239
240    /// Consumes self and returns the underlying map.
241    pub fn into_map(self) -> BTreeMap<StorageMapKey, Word> {
242        self.0
243    }
244}
245
246impl FromIterator<(StorageMapKey, Word)> for StorageMapPatchEntries {
247    /// Creates a new set of map patch entries from the provided iterators of cleared and
248    /// updated entries.
249    fn from_iter<T: IntoIterator<Item = (StorageMapKey, Word)>>(iter: T) -> Self {
250        Self::from_raw(BTreeMap::from_iter(iter))
251    }
252}
253
254/// Converts a [`StorageMap`] into a set of map patch entries for full state patch construction.
255impl From<StorageMap> for StorageMapPatchEntries {
256    fn from(map: StorageMap) -> Self {
257        StorageMapPatchEntries::from_raw(map.into_entries().into_iter().collect())
258    }
259}
260
261impl Serializable for StorageMapPatchEntries {
262    /// Serializes the cleared and updated entries separately. Because the value of a cleared entry
263    /// is always [`Word::empty`], only its key is written, saving [`Word::SERIALIZED_SIZE`] bytes
264    /// per cleared entry compared to writing the empty value.
265    fn write_into<W: ByteWriter>(&self, target: &mut W) {
266        target.write_usize(self.cleared_entries().count());
267        target.write_many(self.cleared_entries());
268
269        target.write_usize(self.updated_entries().count());
270        target.write_many(self.updated_entries());
271    }
272
273    fn get_size_hint(&self) -> usize {
274        let num_cleared = self.cleared_entries().count();
275        let num_updated = self.updated_entries().count();
276
277        num_cleared.get_size_hint()
278            + num_cleared * StorageMapKey::SERIALIZED_SIZE
279            + num_updated.get_size_hint()
280            + num_updated * (StorageMapKey::SERIALIZED_SIZE + Word::SERIALIZED_SIZE)
281    }
282}
283
284impl Deserializable for StorageMapPatchEntries {
285    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
286        let mut entries = BTreeMap::new();
287
288        let num_cleared = source.read_usize()?;
289        for key in source.read_many_iter::<StorageMapKey>(num_cleared)? {
290            let key = key?;
291            if entries.insert(key, Word::empty()).is_some() {
292                return Err(DeserializationError::InvalidValue(format!(
293                    "duplicate key {key} in storage map patch entries"
294                )));
295            }
296        }
297
298        let num_updated = source.read_usize()?;
299        for entry in source.read_many_iter::<(StorageMapKey, Word)>(num_updated)? {
300            let (key, value) = entry?;
301            if entries.insert(key, value).is_some() {
302                return Err(DeserializationError::InvalidValue(format!(
303                    "duplicate key {key} in storage map patch entries"
304                )));
305            }
306        }
307
308        Ok(Self::from_raw(entries))
309    }
310}