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