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