miden-crypto/src/merkle/mmr/partial.rs
src/merkle/mmr/partial.rsRaw
1use super::{MmrDelta, MmrProof, Rpo256, RpoDigest};
2use crate::{
3 merkle::{
4 mmr::{leaf_to_corresponding_tree, nodes_in_forest},
5 InOrderIndex, MerklePath, MmrError, MmrPeaks,
6 },
7 utils::collections::{BTreeMap, Vec},
8};
9
10/// Partially materialized [Mmr], used to efficiently store and update the authentication paths for
11/// a subset of the elements in a full [Mmr].
12///
13/// This structure store only the authentication path for a value, the value itself is stored
14/// separately.
15#[derive(Debug)]
16pub struct PartialMmr {
17 /// The version of the [Mmr].
18 ///
19 /// This value serves the following purposes:
20 ///
21 /// - The forest is a counter for the total number of elements in the [Mmr].
22 /// - Since the [Mmr] is an append-only structure, every change to it causes a change to the
23 /// `forest`, so this value has a dual purpose as a version tag.
24 /// - The bits in the forest also corresponds to the count and size of every perfect binary
25 /// tree that composes the [Mmr] structure, which server to compute indexes and perform
26 /// validation.
27 pub(crate) forest: usize,
28
29 /// The [Mmr] peaks.
30 ///
31 /// The peaks are used for two reasons:
32 ///
33 /// 1. It authenticates the addition of an element to the [PartialMmr], ensuring only valid
34 /// elements are tracked.
35 /// 2. During a [Mmr] update peaks can be merged by hashing the left and right hand sides. The
36 /// peaks are used as the left hand.
37 ///
38 /// All the peaks of every tree in the [Mmr] forest. The peaks are always ordered by number of
39 /// leaves, starting from the peak with most children, to the one with least.
40 pub(crate) peaks: Vec<RpoDigest>,
41
42 /// Authentication nodes used to construct merkle paths for a subset of the [Mmr]'s leaves.
43 ///
44 /// This does not include the [Mmr]'s peaks nor the tracked nodes, only the elements required
45 /// to construct their authentication paths. This property is used to detect when elements can
46 /// be safely removed from, because they are no longer required to authenticate any element in
47 /// the [PartialMmr].
48 ///
49 /// The elements in the [Mmr] are referenced using a in-order tree index. This indexing scheme
50 /// permits for easy computation of the relative nodes (left/right children, sibling, parent),
51 /// which is useful for traversal. The indexing is also stable, meaning that merges to the
52 /// trees in the [Mmr] can be represented without rewrites of the indexes.
53 pub(crate) nodes: BTreeMap<InOrderIndex, RpoDigest>,
54
55 /// Flag indicating if the odd element should be tracked.
56 ///
57 /// This flag is necessary because the sibling of the odd doesn't exist yet, so it can not be
58 /// added into `nodes` to signal the value is being tracked.
59 pub(crate) track_latest: bool,
60}
61
62impl PartialMmr {
63 // CONSTRUCTORS
64 // --------------------------------------------------------------------------------------------
65
66 /// Constructs a [PartialMmr] from the given [MmrPeaks].
67 pub fn from_peaks(accumulator: MmrPeaks) -> Self {
68 let forest = accumulator.num_leaves();
69 let peaks = accumulator.peaks().to_vec();
70 let nodes = BTreeMap::new();
71 let track_latest = false;
72
73 Self { forest, peaks, nodes, track_latest }
74 }
75
76 // ACCESSORS
77 // --------------------------------------------------------------------------------------------
78
79 // Gets the current `forest`.
80 //
81 // This value corresponds to the version of the [PartialMmr] and the number of leaves in it.
82 pub fn forest(&self) -> usize {
83 self.forest
84 }
85
86 // Returns a reference to the current peaks in the [PartialMmr]
87 pub fn peaks(&self) -> &[RpoDigest] {
88 &self.peaks
89 }
90
91 /// Given a leaf position, returns the Merkle path to its corresponding peak. If the position
92 /// is greater-or-equal than the tree size an error is returned. If the requested value is not
93 /// tracked returns `None`.
94 ///
95 /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were
96 /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element
97 /// has position 0, the second position 1, and so on.
98 pub fn open(&self, pos: usize) -> Result<Option<MmrProof>, MmrError> {
99 let tree_bit =
100 leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?;
101 let depth = tree_bit as usize;
102
103 let mut nodes = Vec::with_capacity(depth);
104 let mut idx = InOrderIndex::from_leaf_pos(pos);
105
106 while let Some(node) = self.nodes.get(&idx.sibling()) {
107 nodes.push(*node);
108 idx = idx.parent();
109 }
110
111 // If there are nodes then the path must be complete, otherwise it is a bug
112 debug_assert!(nodes.is_empty() || nodes.len() == depth);
113
114 if nodes.len() != depth {
115 // The requested `pos` is not being tracked.
116 Ok(None)
117 } else {
118 Ok(Some(MmrProof {
119 forest: self.forest,
120 position: pos,
121 merkle_path: MerklePath::new(nodes),
122 }))
123 }
124 }
125
126 // MODIFIERS
127 // --------------------------------------------------------------------------------------------
128
129 /// Add the authentication path represented by [MerklePath] if it is valid.
130 ///
131 /// The `index` refers to the global position of the leaf in the [Mmr], these are 0-indexed
132 /// values assigned in a strictly monotonic fashion as elements are inserted into the [Mmr],
133 /// this value corresponds to the values used in the [Mmr] structure.
134 ///
135 /// The `node` corresponds to the value at `index`, and `path` is the authentication path for
136 /// that element up to its corresponding Mmr peak. The `node` is only used to compute the root
137 /// from the authentication path to valid the data, only the authentication data is saved in
138 /// the structure. If the value is required it should be stored out-of-band.
139 pub fn add(
140 &mut self,
141 index: usize,
142 node: RpoDigest,
143 path: &MerklePath,
144 ) -> Result<(), MmrError> {
145 // Checks there is a tree with same depth as the authentication path, if not the path is
146 // invalid.
147 let tree = 1 << path.depth();
148 if tree & self.forest == 0 {
149 return Err(MmrError::UnknownPeak);
150 };
151
152 if index + 1 == self.forest
153 && path.depth() == 0
154 && self.peaks.last().map_or(false, |v| *v == node)
155 {
156 self.track_latest = true;
157 return Ok(());
158 }
159
160 // ignore the trees smaller than the target (these elements are position after the current
161 // target and don't affect the target index)
162 let target_forest = self.forest ^ (self.forest & (tree - 1));
163 let peak_pos = (target_forest.count_ones() - 1) as usize;
164
165 // translate from mmr index to merkle path
166 let path_idx = index - (target_forest ^ tree);
167
168 // Compute the root of the authentication path, and check it matches the current version of
169 // the PartialMmr.
170 let computed = path.compute_root(path_idx as u64, node).map_err(MmrError::MerkleError)?;
171 if self.peaks[peak_pos] != computed {
172 return Err(MmrError::InvalidPeak);
173 }
174
175 let mut idx = InOrderIndex::from_leaf_pos(index);
176 for node in path.nodes() {
177 self.nodes.insert(idx.sibling(), *node);
178 idx = idx.parent();
179 }
180
181 Ok(())
182 }
183
184 /// Remove a leaf of the [PartialMmr] and the unused nodes from the authentication path.
185 ///
186 /// Note: `leaf_pos` corresponds to the position the [Mmr] and not on an individual tree.
187 pub fn remove(&mut self, leaf_pos: usize) {
188 let mut idx = InOrderIndex::from_leaf_pos(leaf_pos);
189
190 self.nodes.remove(&idx.sibling());
191
192 // `idx` represent the element that can be computed by the authentication path, because
193 // these elements can be computed they are not saved for the authentication of the current
194 // target. In other words, if the idx is present it was added for the authentication of
195 // another element, and no more elements should be removed otherwise it would remove that
196 // element's authentication data.
197 while !self.nodes.contains_key(&idx) {
198 idx = idx.parent();
199 self.nodes.remove(&idx.sibling());
200 }
201 }
202
203 /// Applies updates to the [PartialMmr].
204 pub fn apply(&mut self, delta: MmrDelta) -> Result<(), MmrError> {
205 if delta.forest < self.forest {
206 return Err(MmrError::InvalidPeaks);
207 }
208
209 if delta.forest == self.forest {
210 if !delta.data.is_empty() {
211 return Err(MmrError::InvalidUpdate);
212 }
213
214 return Ok(());
215 }
216
217 // find the tree merges
218 let changes = self.forest ^ delta.forest;
219 let largest = 1 << changes.ilog2();
220 let merges = self.forest & (largest - 1);
221
222 debug_assert!(
223 !self.track_latest || (merges & 1) == 1,
224 "if there is an odd element, a merge is required"
225 );
226
227 // count the number elements needed to produce largest from the current state
228 let (merge_count, new_peaks) = if merges != 0 {
229 let depth = largest.trailing_zeros();
230 let skipped = merges.trailing_zeros();
231 let computed = merges.count_ones() - 1;
232 let merge_count = depth - skipped - computed;
233
234 let new_peaks = delta.forest & (largest - 1);
235
236 (merge_count, new_peaks)
237 } else {
238 (0, changes)
239 };
240
241 // verify the delta size
242 if (delta.data.len() as u32) != merge_count + new_peaks.count_ones() {
243 return Err(MmrError::InvalidUpdate);
244 }
245
246 // keeps track of how many data elements from the update have been consumed
247 let mut update_count = 0;
248
249 if merges != 0 {
250 // starts at the smallest peak and follows the merged peaks
251 let mut peak_idx = forest_to_root_index(self.forest);
252
253 // match order of the update data while applying it
254 self.peaks.reverse();
255
256 // set to true when the data is needed for authentication paths updates
257 let mut track = self.track_latest;
258 self.track_latest = false;
259
260 let mut peak_count = 0;
261 let mut target = 1 << merges.trailing_zeros();
262 let mut new = delta.data[0];
263 update_count += 1;
264
265 while target < largest {
266 // check if either the left or right subtrees have saved for authentication paths.
267 // If so, turn tracking on to update those paths.
268 if target != 1 && !track {
269 let left_child = peak_idx.left_child();
270 let right_child = peak_idx.right_child();
271 track = self.nodes.contains_key(&left_child)
272 | self.nodes.contains_key(&right_child);
273 }
274
275 // update data only contains the nodes from the right subtrees, left nodes are
276 // either previously known peaks or computed values
277 let (left, right) = if target & merges != 0 {
278 let peak = self.peaks[peak_count];
279 peak_count += 1;
280 (peak, new)
281 } else {
282 let update = delta.data[update_count];
283 update_count += 1;
284 (new, update)
285 };
286
287 if track {
288 self.nodes.insert(peak_idx.sibling(), right);
289 }
290
291 peak_idx = peak_idx.parent();
292 new = Rpo256::merge(&[left, right]);
293 target <<= 1;
294 }
295
296 debug_assert!(peak_count == (merges.count_ones() as usize));
297
298 // restore the peaks order
299 self.peaks.reverse();
300 // remove the merged peaks
301 self.peaks.truncate(self.peaks.len() - peak_count);
302 // add the newly computed peak, the result of the merges
303 self.peaks.push(new);
304 }
305
306 // The rest of the update data is composed of peaks. None of these elements can contain
307 // tracked elements because the peaks were unknown, and it is not possible to add elements
308 // for tacking without authenticating it to a peak.
309 self.peaks.extend_from_slice(&delta.data[update_count..]);
310 self.forest = delta.forest;
311
312 debug_assert!(self.peaks.len() == (self.forest.count_ones() as usize));
313
314 Ok(())
315 }
316}
317
318// CONVERSIONS
319// ================================================================================================
320
321impl From<MmrPeaks> for PartialMmr {
322 fn from(peaks: MmrPeaks) -> Self {
323 Self::from_peaks(peaks)
324 }
325}
326
327impl From<PartialMmr> for MmrPeaks {
328 fn from(partial_mmr: PartialMmr) -> Self {
329 // Safety: the [PartialMmr] maintains the constraints the number of true bits in the forest
330 // matches the number of peaks, as required by the [MmrPeaks]
331 MmrPeaks::new(partial_mmr.forest, partial_mmr.peaks).unwrap()
332 }
333}
334
335impl From<&MmrPeaks> for PartialMmr {
336 fn from(peaks: &MmrPeaks) -> Self {
337 Self::from_peaks(peaks.clone())
338 }
339}
340
341impl From<&PartialMmr> for MmrPeaks {
342 fn from(partial_mmr: &PartialMmr) -> Self {
343 // Safety: the [PartialMmr] maintains the constraints the number of true bits in the forest
344 // matches the number of peaks, as required by the [MmrPeaks]
345 MmrPeaks::new(partial_mmr.forest, partial_mmr.peaks.clone()).unwrap()
346 }
347}
348
349// UTILS
350// ================================================================================================
351
352/// Given the description of a `forest`, returns the index of the root element of the smallest tree
353/// in it.
354pub fn forest_to_root_index(forest: usize) -> InOrderIndex {
355 // Count total size of all trees in the forest.
356 let nodes = nodes_in_forest(forest);
357
358 // Add the count for the parent nodes that separate each tree. These are allocated but
359 // currently empty, and correspond to the nodes that will be used once the trees are merged.
360 let open_trees = (forest.count_ones() - 1) as usize;
361
362 // Remove the count of the right subtree of the target tree, target tree root index comes
363 // before the subtree for the in-order tree walk.
364 let right_subtree_count = ((1u32 << forest.trailing_zeros()) - 1) as usize;
365
366 let idx = nodes + open_trees - right_subtree_count;
367
368 InOrderIndex::new(idx.try_into().unwrap())
369}
370
371#[cfg(test)]
372mod test {
373 use super::forest_to_root_index;
374 use crate::merkle::InOrderIndex;
375
376 #[test]
377 fn test_forest_to_root_index() {
378 fn idx(pos: usize) -> InOrderIndex {
379 InOrderIndex::new(pos.try_into().unwrap())
380 }
381
382 // When there is a single tree in the forest, the index is equivalent to the number of
383 // leaves in that tree, which is `2^n`.
384 assert_eq!(forest_to_root_index(0b0001), idx(1));
385 assert_eq!(forest_to_root_index(0b0010), idx(2));
386 assert_eq!(forest_to_root_index(0b0100), idx(4));
387 assert_eq!(forest_to_root_index(0b1000), idx(8));
388
389 assert_eq!(forest_to_root_index(0b0011), idx(5));
390 assert_eq!(forest_to_root_index(0b0101), idx(9));
391 assert_eq!(forest_to_root_index(0b1001), idx(17));
392 assert_eq!(forest_to_root_index(0b0111), idx(13));
393 assert_eq!(forest_to_root_index(0b1011), idx(21));
394 assert_eq!(forest_to_root_index(0b1111), idx(29));
395
396 assert_eq!(forest_to_root_index(0b0110), idx(10));
397 assert_eq!(forest_to_root_index(0b1010), idx(18));
398 assert_eq!(forest_to_root_index(0b1100), idx(20));
399 assert_eq!(forest_to_root_index(0b1110), idx(26));
400 }
401}