miden-crypto/src/merkle/partial_mt/mod.rs
src/merkle/partial_mt/mod.rsRaw
1use super::{
2 BTreeMap, BTreeSet, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, ValuePath, Vec,
3 Word, EMPTY_WORD,
4};
5use crate::utils::{format, string::String, word_to_hex};
6use core::fmt;
7
8#[cfg(test)]
9mod tests;
10
11// CONSTANTS
12// ================================================================================================
13
14/// Index of the root node.
15const ROOT_INDEX: NodeIndex = NodeIndex::root();
16
17/// An RpoDigest consisting of 4 ZERO elements.
18const EMPTY_DIGEST: RpoDigest = RpoDigest::new(EMPTY_WORD);
19
20// PARTIAL MERKLE TREE
21// ================================================================================================
22
23/// A partial Merkle tree with NodeIndex keys and 4-element RpoDigest leaf values. Partial Merkle
24/// Tree allows to create Merkle Tree by providing Merkle paths of different lengths.
25///
26/// The root of the tree is recomputed on each new leaf update.
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct PartialMerkleTree {
29 max_depth: u8,
30 nodes: BTreeMap<NodeIndex, RpoDigest>,
31 leaves: BTreeSet<NodeIndex>,
32}
33
34impl Default for PartialMerkleTree {
35 fn default() -> Self {
36 Self::new()
37 }
38}
39
40impl PartialMerkleTree {
41 // CONSTANTS
42 // --------------------------------------------------------------------------------------------
43
44 /// Minimum supported depth.
45 pub const MIN_DEPTH: u8 = 1;
46
47 /// Maximum supported depth.
48 pub const MAX_DEPTH: u8 = 64;
49
50 // CONSTRUCTORS
51 // --------------------------------------------------------------------------------------------
52
53 /// Returns a new emply [PartialMerkleTree].
54 pub fn new() -> Self {
55 PartialMerkleTree {
56 max_depth: 0,
57 nodes: BTreeMap::new(),
58 leaves: BTreeSet::new(),
59 }
60 }
61
62 /// Appends the provided paths iterator into the set.
63 ///
64 /// Analogous to [Self::add_path].
65 pub fn with_paths<I>(paths: I) -> Result<Self, MerkleError>
66 where
67 I: IntoIterator<Item = (u64, RpoDigest, MerklePath)>,
68 {
69 // create an empty tree
70 let tree = PartialMerkleTree::new();
71
72 paths.into_iter().try_fold(tree, |mut tree, (index, value, path)| {
73 tree.add_path(index, value, path)?;
74 Ok(tree)
75 })
76 }
77
78 // PUBLIC ACCESSORS
79 // --------------------------------------------------------------------------------------------
80
81 /// Returns the root of this Merkle tree.
82 pub fn root(&self) -> RpoDigest {
83 self.nodes.get(&ROOT_INDEX).cloned().unwrap_or(EMPTY_DIGEST)
84 }
85
86 /// Returns the depth of this Merkle tree.
87 pub fn max_depth(&self) -> u8 {
88 self.max_depth
89 }
90
91 /// Returns a node at the specified NodeIndex.
92 ///
93 /// # Errors
94 /// Returns an error if the specified NodeIndex is not contained in the nodes map.
95 pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
96 self.nodes.get(&index).ok_or(MerkleError::NodeNotInSet(index)).map(|hash| *hash)
97 }
98
99 /// Returns true if provided index contains in the leaves set, false otherwise.
100 pub fn is_leaf(&self, index: NodeIndex) -> bool {
101 self.leaves.contains(&index)
102 }
103
104 /// Returns a vector of paths from every leaf to the root.
105 pub fn paths(&self) -> Vec<(NodeIndex, ValuePath)> {
106 let mut paths = Vec::new();
107 self.leaves.iter().for_each(|&leaf| {
108 paths.push((
109 leaf,
110 ValuePath {
111 value: *self.get_node(leaf).expect("Failed to get leaf node"),
112 path: self.get_path(leaf).expect("Failed to get path"),
113 },
114 ));
115 });
116 paths
117 }
118
119 /// Returns a Merkle path from the node at the specified index to the root.
120 ///
121 /// The node itself is not included in the path.
122 ///
123 /// # Errors
124 /// Returns an error if:
125 /// - the specified index has depth set to 0 or the depth is greater than the depth of this
126 /// Merkle tree.
127 /// - the specified index is not contained in the nodes map.
128 pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
129 if index.is_root() {
130 return Err(MerkleError::DepthTooSmall(index.depth()));
131 } else if index.depth() > self.max_depth() {
132 return Err(MerkleError::DepthTooBig(index.depth() as u64));
133 }
134
135 if !self.nodes.contains_key(&index) {
136 return Err(MerkleError::NodeNotInSet(index));
137 }
138
139 let mut path = Vec::new();
140 for _ in 0..index.depth() {
141 let sibling_index = index.sibling();
142 index.move_up();
143 let sibling =
144 self.nodes.get(&sibling_index).cloned().expect("Sibling node not in the map");
145 path.push(Word::from(sibling));
146 }
147 Ok(MerklePath::new(path))
148 }
149
150 // ITERATORS
151 // --------------------------------------------------------------------------------------------
152
153 /// Returns an iterator over the leaves of this [PartialMerkleTree].
154 pub fn leaves(&self) -> impl Iterator<Item = (NodeIndex, RpoDigest)> + '_ {
155 self.leaves.iter().map(|&leaf| {
156 (
157 leaf,
158 self.get_node(leaf).expect(&format!("Leaf with {leaf} is not in the nodes map")),
159 )
160 })
161 }
162
163 // STATE MUTATORS
164 // --------------------------------------------------------------------------------------------
165
166 /// Adds the nodes of the specified Merkle path to this [PartialMerkleTree]. The `index_value`
167 /// and `value` parameters specify the leaf node at which the path starts.
168 ///
169 /// # Errors
170 /// Returns an error if:
171 /// - The depth of the specified node_index is greater than 64 or smaller than 1.
172 /// - The specified path is not consistent with other paths in the set (i.e., resolves to a
173 /// different root).
174 pub fn add_path(
175 &mut self,
176 index_value: u64,
177 value: RpoDigest,
178 path: MerklePath,
179 ) -> Result<(), MerkleError> {
180 let index_value = NodeIndex::new(path.len() as u8, index_value)?;
181
182 Self::check_depth(index_value.depth())?;
183 self.update_depth(index_value.depth());
184
185 // add provided node and its sibling to the leaves set
186 self.leaves.insert(index_value);
187 let sibling_node_index = index_value.sibling();
188 self.leaves.insert(sibling_node_index);
189
190 // add provided node and its sibling to the nodes map
191 self.nodes.insert(index_value, value);
192 self.nodes.insert(sibling_node_index, path[0].into());
193
194 // traverse to the root, updating the nodes
195 let mut index_value = index_value;
196 let node = Rpo256::merge(&index_value.build_node(value, path[0].into()));
197 let root = path.iter().skip(1).copied().fold(node, |node, hash| {
198 index_value.move_up();
199 // insert calculated node to the nodes map
200 self.nodes.insert(index_value, node);
201
202 // if the calculated node was a leaf, remove it from leaves set.
203 self.leaves.remove(&index_value);
204
205 let sibling_node = index_value.sibling();
206
207 // Insert node from Merkle path to the nodes map. This sibling node becomes a leaf only
208 // if it is a new node (it wasn't in nodes map).
209 // Node can be in 3 states: internal node, leaf of the tree and not a tree node at all.
210 // - Internal node can only stay in this state -- addition of a new path can't make it
211 // a leaf or remove it from the tree.
212 // - Leaf node can stay in the same state (remain a leaf) or can become an internal
213 // node. In the first case we don't need to do anything, and the second case is handled
214 // by the call of `self.leaves.remove(&index_value);`
215 // - New node can be a calculated node or a "sibling" node from a Merkle Path:
216 // --- Calculated node, obviously, never can be a leaf.
217 // --- Sibling node can be only a leaf, because otherwise it is not a new node.
218 if self.nodes.insert(sibling_node, hash.into()).is_none() {
219 self.leaves.insert(sibling_node);
220 }
221
222 Rpo256::merge(&index_value.build_node(node, hash.into()))
223 });
224
225 // if the path set is empty (the root is all ZEROs), set the root to the root of the added
226 // path; otherwise, the root of the added path must be identical to the current root
227 if self.root() == EMPTY_DIGEST {
228 self.nodes.insert(ROOT_INDEX, root);
229 } else if self.root() != root {
230 return Err(MerkleError::ConflictingRoots([*self.root(), *root].to_vec()));
231 }
232
233 Ok(())
234 }
235
236 /// Updates value of the leaf at the specified index returning the old leaf value.
237 ///
238 /// This also recomputes all hashes between the leaf and the root, updating the root itself.
239 ///
240 /// # Errors
241 /// Returns an error if:
242 /// - The depth of the specified node_index is greater than 64 or smaller than 1.
243 /// - The specified node index is not corresponding to the leaf.
244 pub fn update_leaf(
245 &mut self,
246 node_index: NodeIndex,
247 value: RpoDigest,
248 ) -> Result<RpoDigest, MerkleError> {
249 // check correctness of the depth and update it
250 Self::check_depth(node_index.depth())?;
251 self.update_depth(node_index.depth());
252
253 // insert NodeIndex to the leaves Set
254 self.leaves.insert(node_index);
255
256 // add node value to the nodes Map
257 let old_value = self
258 .nodes
259 .insert(node_index, value)
260 .ok_or(MerkleError::NodeNotInSet(node_index))?;
261
262 // if the old value and new value are the same, there is nothing to update
263 if value == old_value {
264 return Ok(old_value);
265 }
266
267 let mut node_index = node_index;
268 let mut value = value;
269 for _ in 0..node_index.depth() {
270 let sibling = self.nodes.get(&node_index.sibling()).expect("sibling should exist");
271 value = Rpo256::merge(&node_index.build_node(value, *sibling));
272 node_index.move_up();
273 self.nodes.insert(node_index, value);
274 }
275
276 Ok(old_value)
277 }
278
279 // UTILITY FUNCTIONS
280 // --------------------------------------------------------------------------------------------
281
282 /// Utility to visualize a [PartialMerkleTree] in text.
283 pub fn print(&self) -> Result<String, fmt::Error> {
284 let indent = " ";
285 let mut s = String::new();
286 s.push_str("root: ");
287 s.push_str(&word_to_hex(&self.root())?);
288 s.push('\n');
289 for d in 1..=self.max_depth() {
290 let entries = 2u64.pow(d.into());
291 for i in 0..entries {
292 let index = NodeIndex::new(d, i).expect("The index must always be valid");
293 let node = self.get_node(index);
294 let node = match node {
295 Err(_) => continue,
296 Ok(node) => node,
297 };
298
299 for _ in 0..d {
300 s.push_str(indent);
301 }
302 s.push_str(&format!("({}, {}): ", index.depth(), index.value()));
303 s.push_str(&word_to_hex(&node)?);
304 s.push('\n');
305 }
306 }
307
308 Ok(s)
309 }
310
311 // HELPER METHODS
312 // --------------------------------------------------------------------------------------------
313
314 /// Updates depth value with the maximum of current and provided depth.
315 fn update_depth(&mut self, new_depth: u8) {
316 self.max_depth = new_depth.max(self.max_depth);
317 }
318
319 /// Returns an error if the depth is 0 or is greater than 64.
320 fn check_depth(depth: u8) -> Result<(), MerkleError> {
321 // validate the range of the depth.
322 if depth < Self::MIN_DEPTH {
323 return Err(MerkleError::DepthTooSmall(depth));
324 } else if Self::MAX_DEPTH < depth {
325 return Err(MerkleError::DepthTooBig(depth as u64));
326 }
327 Ok(())
328 }
329}