miden-crypto/src/merkle/partial_mt/mod.rs
src/merkle/partial_mt/mod.rsRaw
1use core::fmt;
2
3use super::{
4 InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, ValuePath, Word,
5 EMPTY_WORD,
6};
7use crate::utils::{
8 collections::*, format, string::*, vec, word_to_hex, ByteReader, ByteWriter, Deserializable,
9 DeserializationError, Serializable,
10};
11
12#[cfg(test)]
13mod tests;
14
15// CONSTANTS
16// ================================================================================================
17
18/// Index of the root node.
19const ROOT_INDEX: NodeIndex = NodeIndex::root();
20
21/// An RpoDigest consisting of 4 ZERO elements.
22const EMPTY_DIGEST: RpoDigest = RpoDigest::new(EMPTY_WORD);
23
24// PARTIAL MERKLE TREE
25// ================================================================================================
26
27/// A partial Merkle tree with NodeIndex keys and 4-element RpoDigest leaf values. Partial Merkle
28/// Tree allows to create Merkle Tree by providing Merkle paths of different lengths.
29///
30/// The root of the tree is recomputed on each new leaf update.
31#[derive(Debug, Clone, PartialEq, Eq)]
32#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
33pub struct PartialMerkleTree {
34 max_depth: u8,
35 nodes: BTreeMap<NodeIndex, RpoDigest>,
36 leaves: BTreeSet<NodeIndex>,
37}
38
39impl Default for PartialMerkleTree {
40 fn default() -> Self {
41 Self::new()
42 }
43}
44
45impl PartialMerkleTree {
46 // CONSTANTS
47 // --------------------------------------------------------------------------------------------
48
49 /// Minimum supported depth.
50 pub const MIN_DEPTH: u8 = 1;
51
52 /// Maximum supported depth.
53 pub const MAX_DEPTH: u8 = 64;
54
55 // CONSTRUCTORS
56 // --------------------------------------------------------------------------------------------
57
58 /// Returns a new empty [PartialMerkleTree].
59 pub fn new() -> Self {
60 PartialMerkleTree {
61 max_depth: 0,
62 nodes: BTreeMap::new(),
63 leaves: BTreeSet::new(),
64 }
65 }
66
67 /// Appends the provided paths iterator into the set.
68 ///
69 /// Analogous to [Self::add_path].
70 pub fn with_paths<I>(paths: I) -> Result<Self, MerkleError>
71 where
72 I: IntoIterator<Item = (u64, RpoDigest, MerklePath)>,
73 {
74 // create an empty tree
75 let tree = PartialMerkleTree::new();
76
77 paths.into_iter().try_fold(tree, |mut tree, (index, value, path)| {
78 tree.add_path(index, value, path)?;
79 Ok(tree)
80 })
81 }
82
83 /// Returns a new [PartialMerkleTree] instantiated with leaves map as specified by the provided
84 /// entries.
85 ///
86 /// # Errors
87 /// Returns an error if:
88 /// - If the depth is 0 or is greater than 64.
89 /// - The number of entries exceeds the maximum tree capacity, that is 2^{depth}.
90 /// - The provided entries contain an insufficient set of nodes.
91 pub fn with_leaves<R, I>(entries: R) -> Result<Self, MerkleError>
92 where
93 R: IntoIterator<IntoIter = I>,
94 I: Iterator<Item = (NodeIndex, RpoDigest)> + ExactSizeIterator,
95 {
96 let mut layers: BTreeMap<u8, Vec<u64>> = BTreeMap::new();
97 let mut leaves = BTreeSet::new();
98 let mut nodes = BTreeMap::new();
99
100 // add data to the leaves and nodes maps and also fill layers map, where the key is the
101 // depth of the node and value is its index.
102 for (node_index, hash) in entries.into_iter() {
103 leaves.insert(node_index);
104 nodes.insert(node_index, hash);
105 layers
106 .entry(node_index.depth())
107 .and_modify(|layer_vec| layer_vec.push(node_index.value()))
108 .or_insert(vec![node_index.value()]);
109 }
110
111 // check if the number of leaves can be accommodated by the tree's depth; we use a min
112 // depth of 63 because we consider passing in a vector of size 2^64 infeasible.
113 let max = 2usize.pow(63);
114 if layers.len() > max {
115 return Err(MerkleError::InvalidNumEntries(max));
116 }
117
118 // Get maximum depth
119 let max_depth = *layers.keys().next_back().unwrap_or(&0);
120
121 // fill layers without nodes with empty vector
122 for depth in 0..max_depth {
123 layers.entry(depth).or_default();
124 }
125
126 let mut layer_iter = layers.into_values().rev();
127 let mut parent_layer = layer_iter.next().unwrap();
128 let mut current_layer;
129
130 for depth in (1..max_depth + 1).rev() {
131 // set current_layer = parent_layer and parent_layer = layer_iter.next()
132 current_layer = layer_iter.next().unwrap();
133 core::mem::swap(&mut current_layer, &mut parent_layer);
134
135 for index_value in current_layer {
136 // get the parent node index
137 let parent_node = NodeIndex::new(depth - 1, index_value / 2)?;
138
139 // Check if the parent hash was already calculated. In about half of the cases, we
140 // don't need to do anything.
141 if !parent_layer.contains(&parent_node.value()) {
142 // create current node index
143 let index = NodeIndex::new(depth, index_value)?;
144
145 // get hash of the current node
146 let node = nodes.get(&index).ok_or(MerkleError::NodeNotInSet(index))?;
147 // get hash of the sibling node
148 let sibling = nodes
149 .get(&index.sibling())
150 .ok_or(MerkleError::NodeNotInSet(index.sibling()))?;
151 // get parent hash
152 let parent = Rpo256::merge(&index.build_node(*node, *sibling));
153
154 // add index value of the calculated node to the parents layer
155 parent_layer.push(parent_node.value());
156 // add index and hash to the nodes map
157 nodes.insert(parent_node, parent);
158 }
159 }
160 }
161
162 Ok(PartialMerkleTree { max_depth, nodes, leaves })
163 }
164
165 // PUBLIC ACCESSORS
166 // --------------------------------------------------------------------------------------------
167
168 /// Returns the root of this Merkle tree.
169 pub fn root(&self) -> RpoDigest {
170 self.nodes.get(&ROOT_INDEX).cloned().unwrap_or(EMPTY_DIGEST)
171 }
172
173 /// Returns the depth of this Merkle tree.
174 pub fn max_depth(&self) -> u8 {
175 self.max_depth
176 }
177
178 /// Returns a node at the specified NodeIndex.
179 ///
180 /// # Errors
181 /// Returns an error if the specified NodeIndex is not contained in the nodes map.
182 pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
183 self.nodes.get(&index).ok_or(MerkleError::NodeNotInSet(index)).copied()
184 }
185
186 /// Returns true if provided index contains in the leaves set, false otherwise.
187 pub fn is_leaf(&self, index: NodeIndex) -> bool {
188 self.leaves.contains(&index)
189 }
190
191 /// Returns a vector of paths from every leaf to the root.
192 pub fn to_paths(&self) -> Vec<(NodeIndex, ValuePath)> {
193 let mut paths = Vec::new();
194 self.leaves.iter().for_each(|&leaf| {
195 paths.push((
196 leaf,
197 ValuePath {
198 value: self.get_node(leaf).expect("Failed to get leaf node"),
199 path: self.get_path(leaf).expect("Failed to get path"),
200 },
201 ));
202 });
203 paths
204 }
205
206 /// Returns a Merkle path from the node at the specified index to the root.
207 ///
208 /// The node itself is not included in the path.
209 ///
210 /// # Errors
211 /// Returns an error if:
212 /// - the specified index has depth set to 0 or the depth is greater than the depth of this
213 /// Merkle tree.
214 /// - the specified index is not contained in the nodes map.
215 pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
216 if index.is_root() {
217 return Err(MerkleError::DepthTooSmall(index.depth()));
218 } else if index.depth() > self.max_depth() {
219 return Err(MerkleError::DepthTooBig(index.depth() as u64));
220 }
221
222 if !self.nodes.contains_key(&index) {
223 return Err(MerkleError::NodeNotInSet(index));
224 }
225
226 let mut path = Vec::new();
227 for _ in 0..index.depth() {
228 let sibling_index = index.sibling();
229 index.move_up();
230 let sibling =
231 self.nodes.get(&sibling_index).cloned().expect("Sibling node not in the map");
232 path.push(sibling);
233 }
234 Ok(MerklePath::new(path))
235 }
236
237 // ITERATORS
238 // --------------------------------------------------------------------------------------------
239
240 /// Returns an iterator over the leaves of this [PartialMerkleTree].
241 pub fn leaves(&self) -> impl Iterator<Item = (NodeIndex, RpoDigest)> + '_ {
242 self.leaves.iter().map(|&leaf| {
243 (
244 leaf,
245 self.get_node(leaf)
246 .unwrap_or_else(|_| panic!("Leaf with {leaf} is not in the nodes map")),
247 )
248 })
249 }
250
251 /// Returns an iterator over the inner nodes of this Merkle tree.
252 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
253 let inner_nodes = self.nodes.iter().filter(|(index, _)| !self.leaves.contains(index));
254 inner_nodes.map(|(index, digest)| {
255 let left_hash =
256 self.nodes.get(&index.left_child()).expect("Failed to get left child hash");
257 let right_hash =
258 self.nodes.get(&index.right_child()).expect("Failed to get right child hash");
259 InnerNodeInfo {
260 value: *digest,
261 left: *left_hash,
262 right: *right_hash,
263 }
264 })
265 }
266
267 // STATE MUTATORS
268 // --------------------------------------------------------------------------------------------
269
270 /// Adds the nodes of the specified Merkle path to this [PartialMerkleTree]. The `index_value`
271 /// and `value` parameters specify the leaf node at which the path starts.
272 ///
273 /// # Errors
274 /// Returns an error if:
275 /// - The depth of the specified node_index is greater than 64 or smaller than 1.
276 /// - The specified path is not consistent with other paths in the set (i.e., resolves to a
277 /// different root).
278 pub fn add_path(
279 &mut self,
280 index_value: u64,
281 value: RpoDigest,
282 path: MerklePath,
283 ) -> Result<(), MerkleError> {
284 let index_value = NodeIndex::new(path.len() as u8, index_value)?;
285
286 Self::check_depth(index_value.depth())?;
287 self.update_depth(index_value.depth());
288
289 // add provided node and its sibling to the leaves set
290 self.leaves.insert(index_value);
291 let sibling_node_index = index_value.sibling();
292 self.leaves.insert(sibling_node_index);
293
294 // add provided node and its sibling to the nodes map
295 self.nodes.insert(index_value, value);
296 self.nodes.insert(sibling_node_index, path[0]);
297
298 // traverse to the root, updating the nodes
299 let mut index_value = index_value;
300 let node = Rpo256::merge(&index_value.build_node(value, path[0]));
301 let root = path.iter().skip(1).copied().fold(node, |node, hash| {
302 index_value.move_up();
303 // insert calculated node to the nodes map
304 self.nodes.insert(index_value, node);
305
306 // if the calculated node was a leaf, remove it from leaves set.
307 self.leaves.remove(&index_value);
308
309 let sibling_node = index_value.sibling();
310
311 // Insert node from Merkle path to the nodes map. This sibling node becomes a leaf only
312 // if it is a new node (it wasn't in nodes map).
313 // Node can be in 3 states: internal node, leaf of the tree and not a tree node at all.
314 // - Internal node can only stay in this state -- addition of a new path can't make it
315 // a leaf or remove it from the tree.
316 // - Leaf node can stay in the same state (remain a leaf) or can become an internal
317 // node. In the first case we don't need to do anything, and the second case is handled
318 // by the call of `self.leaves.remove(&index_value);`
319 // - New node can be a calculated node or a "sibling" node from a Merkle Path:
320 // --- Calculated node, obviously, never can be a leaf.
321 // --- Sibling node can be only a leaf, because otherwise it is not a new node.
322 if self.nodes.insert(sibling_node, hash).is_none() {
323 self.leaves.insert(sibling_node);
324 }
325
326 Rpo256::merge(&index_value.build_node(node, hash))
327 });
328
329 // if the path set is empty (the root is all ZEROs), set the root to the root of the added
330 // path; otherwise, the root of the added path must be identical to the current root
331 if self.root() == EMPTY_DIGEST {
332 self.nodes.insert(ROOT_INDEX, root);
333 } else if self.root() != root {
334 return Err(MerkleError::ConflictingRoots([self.root(), root].to_vec()));
335 }
336
337 Ok(())
338 }
339
340 /// Updates value of the leaf at the specified index returning the old leaf value.
341 /// By default the specified index is assumed to belong to the deepest layer. If the considered
342 /// node does not belong to the tree, the first node on the way to the root will be changed.
343 ///
344 /// By default the specified index is assumed to belong to the deepest layer. If the considered
345 /// node does not belong to the tree, the first node on the way to the root will be changed.
346 ///
347 /// This also recomputes all hashes between the leaf and the root, updating the root itself.
348 ///
349 /// # Errors
350 /// Returns an error if:
351 /// - The specified index is greater than the maximum number of nodes on the deepest layer.
352 pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<RpoDigest, MerkleError> {
353 let mut node_index = NodeIndex::new(self.max_depth(), index)?;
354
355 // proceed to the leaf
356 for _ in 0..node_index.depth() {
357 if !self.leaves.contains(&node_index) {
358 node_index.move_up();
359 }
360 }
361
362 // add node value to the nodes Map
363 let old_value = self
364 .nodes
365 .insert(node_index, value.into())
366 .ok_or(MerkleError::NodeNotInSet(node_index))?;
367
368 // if the old value and new value are the same, there is nothing to update
369 if value == *old_value {
370 return Ok(old_value);
371 }
372
373 let mut value = value.into();
374 for _ in 0..node_index.depth() {
375 let sibling = self.nodes.get(&node_index.sibling()).expect("sibling should exist");
376 value = Rpo256::merge(&node_index.build_node(value, *sibling));
377 node_index.move_up();
378 self.nodes.insert(node_index, value);
379 }
380
381 Ok(old_value)
382 }
383
384 // UTILITY FUNCTIONS
385 // --------------------------------------------------------------------------------------------
386
387 /// Utility to visualize a [PartialMerkleTree] in text.
388 pub fn print(&self) -> Result<String, fmt::Error> {
389 let indent = " ";
390 let mut s = String::new();
391 s.push_str("root: ");
392 s.push_str(&word_to_hex(&self.root())?);
393 s.push('\n');
394 for d in 1..=self.max_depth() {
395 let entries = 2u64.pow(d.into());
396 for i in 0..entries {
397 let index = NodeIndex::new(d, i).expect("The index must always be valid");
398 let node = self.get_node(index);
399 let node = match node {
400 Err(_) => continue,
401 Ok(node) => node,
402 };
403
404 for _ in 0..d {
405 s.push_str(indent);
406 }
407 s.push_str(&format!("({}, {}): ", index.depth(), index.value()));
408 s.push_str(&word_to_hex(&node)?);
409 s.push('\n');
410 }
411 }
412
413 Ok(s)
414 }
415
416 // HELPER METHODS
417 // --------------------------------------------------------------------------------------------
418
419 /// Updates depth value with the maximum of current and provided depth.
420 fn update_depth(&mut self, new_depth: u8) {
421 self.max_depth = new_depth.max(self.max_depth);
422 }
423
424 /// Returns an error if the depth is 0 or is greater than 64.
425 fn check_depth(depth: u8) -> Result<(), MerkleError> {
426 // validate the range of the depth.
427 if depth < Self::MIN_DEPTH {
428 return Err(MerkleError::DepthTooSmall(depth));
429 } else if Self::MAX_DEPTH < depth {
430 return Err(MerkleError::DepthTooBig(depth as u64));
431 }
432 Ok(())
433 }
434}
435
436// SERIALIZATION
437// ================================================================================================
438
439impl Serializable for PartialMerkleTree {
440 fn write_into<W: ByteWriter>(&self, target: &mut W) {
441 // write leaf nodes
442 target.write_u64(self.leaves.len() as u64);
443 for leaf_index in self.leaves.iter() {
444 leaf_index.write_into(target);
445 self.get_node(*leaf_index).expect("Leaf hash not found").write_into(target);
446 }
447 }
448}
449
450impl Deserializable for PartialMerkleTree {
451 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
452 let leaves_len = source.read_u64()? as usize;
453 let mut leaf_nodes = Vec::with_capacity(leaves_len);
454
455 // add leaf nodes to the vector
456 for _ in 0..leaves_len {
457 let index = NodeIndex::read_from(source)?;
458 let hash = RpoDigest::read_from(source)?;
459 leaf_nodes.push((index, hash));
460 }
461
462 let pmt = PartialMerkleTree::with_leaves(leaf_nodes).map_err(|_| {
463 DeserializationError::InvalidValue("Invalid data for PartialMerkleTree creation".into())
464 })?;
465
466 Ok(pmt)
467 }
468}