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