miden-crypto/src/merkle/path_set.rs
src/merkle/path_set.rsRaw
1use super::{BTreeMap, MerkleError, MerklePath, NodeIndex, Rpo256, ValuePath, Vec};
2use crate::{hash::rpo::RpoDigest, Word};
3
4// MERKLE PATH SET
5// ================================================================================================
6
7/// A set of Merkle paths.
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub struct MerklePathSet {
10 root: RpoDigest,
11 total_depth: u8,
12 paths: BTreeMap<u64, MerklePath>,
13}
14
15impl MerklePathSet {
16 // CONSTRUCTOR
17 // --------------------------------------------------------------------------------------------
18
19 /// Returns an empty MerklePathSet.
20 pub fn new(depth: u8) -> Self {
21 let root = RpoDigest::default();
22 let paths = BTreeMap::new();
23
24 Self {
25 root,
26 total_depth: depth,
27 paths,
28 }
29 }
30
31 /// Appends the provided paths iterator into the set.
32 ///
33 /// Analogous to `[Self::add_path]`.
34 pub fn with_paths<I>(self, paths: I) -> Result<Self, MerkleError>
35 where
36 I: IntoIterator<Item = (u64, RpoDigest, MerklePath)>,
37 {
38 paths.into_iter().try_fold(self, |mut set, (index, value, path)| {
39 set.add_path(index, value.into(), path)?;
40 Ok(set)
41 })
42 }
43
44 // PUBLIC ACCESSORS
45 // --------------------------------------------------------------------------------------------
46
47 /// Returns the root to which all paths in this set resolve.
48 pub const fn root(&self) -> RpoDigest {
49 self.root
50 }
51
52 /// Returns the depth of the Merkle tree implied by the paths stored in this set.
53 ///
54 /// Merkle tree of depth 1 has two leaves, depth 2 has four leaves etc.
55 pub const fn depth(&self) -> u8 {
56 self.total_depth
57 }
58
59 /// Returns a node at the specified index.
60 ///
61 /// # Errors
62 /// Returns an error if:
63 /// * The specified index is not valid for the depth of structure.
64 /// * Requested node does not exist in the set.
65 pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
66 if index.depth() != self.total_depth {
67 return Err(MerkleError::InvalidDepth {
68 expected: self.total_depth,
69 provided: index.depth(),
70 });
71 }
72
73 let parity = index.value() & 1;
74 let path_key = index.value() - parity;
75 self.paths
76 .get(&path_key)
77 .ok_or(MerkleError::NodeNotInSet(index))
78 .map(|path| path[parity as usize])
79 }
80
81 /// Returns a leaf at the specified index.
82 ///
83 /// # Errors
84 /// * The specified index is not valid for the depth of the structure.
85 /// * Leaf with the requested path does not exist in the set.
86 pub fn get_leaf(&self, index: u64) -> Result<Word, MerkleError> {
87 let index = NodeIndex::new(self.depth(), index)?;
88 Ok(self.get_node(index)?.into())
89 }
90
91 /// Returns a Merkle path to the node at the specified index. The node itself is
92 /// not included in the path.
93 ///
94 /// # Errors
95 /// Returns an error if:
96 /// * The specified index is not valid for the depth of structure.
97 /// * Node of the requested path does not exist in the set.
98 pub fn get_path(&self, index: NodeIndex) -> Result<MerklePath, MerkleError> {
99 if index.depth() != self.total_depth {
100 return Err(MerkleError::InvalidDepth {
101 expected: self.total_depth,
102 provided: index.depth(),
103 });
104 }
105
106 let parity = index.value() & 1;
107 let path_key = index.value() - parity;
108 let mut path =
109 self.paths.get(&path_key).cloned().ok_or(MerkleError::NodeNotInSet(index))?;
110 path.remove(parity as usize);
111 Ok(path)
112 }
113
114 /// Returns all paths in this path set together with their indexes.
115 pub fn to_paths(&self) -> Vec<(u64, ValuePath)> {
116 let mut result = Vec::with_capacity(self.paths.len() * 2);
117
118 for (&index, path) in self.paths.iter() {
119 // push path for the even index into the result
120 let path1 = ValuePath {
121 value: path[0],
122 path: MerklePath::new(path[1..].to_vec()),
123 };
124 result.push((index, path1));
125
126 // push path for the odd index into the result
127 let mut path2 = path.clone();
128 let leaf2 = path2.remove(1);
129 let path2 = ValuePath {
130 value: leaf2,
131 path: path2,
132 };
133 result.push((index + 1, path2));
134 }
135
136 result
137 }
138
139 // STATE MUTATORS
140 // --------------------------------------------------------------------------------------------
141
142 /// Adds the specified Merkle path to this [MerklePathSet]. The `index` and `value` parameters
143 /// specify the leaf node at which the path starts.
144 ///
145 /// # Errors
146 /// Returns an error if:
147 /// - The specified index is is not valid in the context of this Merkle path set (i.e., the
148 /// index implies a greater depth than is specified for this set).
149 /// - The specified path is not consistent with other paths in the set (i.e., resolves to a
150 /// different root).
151 pub fn add_path(
152 &mut self,
153 index_value: u64,
154 value: Word,
155 mut path: MerklePath,
156 ) -> Result<(), MerkleError> {
157 let mut index = NodeIndex::new(path.len() as u8, index_value)?;
158 if index.depth() != self.total_depth {
159 return Err(MerkleError::InvalidDepth {
160 expected: self.total_depth,
161 provided: index.depth(),
162 });
163 }
164
165 // update the current path
166 let parity = index_value & 1;
167 path.insert(parity as usize, value.into());
168
169 // traverse to the root, updating the nodes
170 let root = Rpo256::merge(&[path[0], path[1]]);
171 let root = path.iter().skip(2).copied().fold(root, |root, hash| {
172 index.move_up();
173 Rpo256::merge(&index.build_node(root, hash))
174 });
175
176 // if the path set is empty (the root is all ZEROs), set the root to the root of the added
177 // path; otherwise, the root of the added path must be identical to the current root
178 if self.root == RpoDigest::default() {
179 self.root = root;
180 } else if self.root != root {
181 return Err(MerkleError::ConflictingRoots([self.root, root].to_vec()));
182 }
183
184 // finish updating the path
185 let path_key = index_value - parity;
186 self.paths.insert(path_key, path);
187 Ok(())
188 }
189
190 /// Replaces the leaf at the specified index with the provided value.
191 ///
192 /// # Errors
193 /// Returns an error if:
194 /// * Requested node does not exist in the set.
195 pub fn update_leaf(&mut self, base_index_value: u64, value: Word) -> Result<(), MerkleError> {
196 let mut index = NodeIndex::new(self.depth(), base_index_value)?;
197 let parity = index.value() & 1;
198 let path_key = index.value() - parity;
199 let path = match self.paths.get_mut(&path_key) {
200 Some(path) => path,
201 None => return Err(MerkleError::NodeNotInSet(index)),
202 };
203
204 // Fill old_hashes vector -----------------------------------------------------------------
205 let mut current_index = index;
206 let mut old_hashes = Vec::with_capacity(path.len().saturating_sub(2));
207 let mut root = Rpo256::merge(&[path[0], path[1]]);
208 for hash in path.iter().skip(2).copied() {
209 old_hashes.push(root);
210 current_index.move_up();
211 let input = current_index.build_node(hash, root);
212 root = Rpo256::merge(&input);
213 }
214
215 // Fill new_hashes vector -----------------------------------------------------------------
216 path[index.is_value_odd() as usize] = value.into();
217
218 let mut new_hashes = Vec::with_capacity(path.len().saturating_sub(2));
219 let mut new_root = Rpo256::merge(&[path[0], path[1]]);
220 for path_hash in path.iter().skip(2).copied() {
221 new_hashes.push(new_root);
222 index.move_up();
223 let input = current_index.build_node(path_hash, new_root);
224 new_root = Rpo256::merge(&input);
225 }
226
227 self.root = new_root;
228
229 // update paths ---------------------------------------------------------------------------
230 for path in self.paths.values_mut() {
231 for i in (0..old_hashes.len()).rev() {
232 if path[i + 2] == old_hashes[i] {
233 path[i + 2] = new_hashes[i];
234 break;
235 }
236 }
237 }
238
239 Ok(())
240 }
241}
242
243// TESTS
244// ================================================================================================
245
246#[cfg(test)]
247mod tests {
248 use super::*;
249 use crate::merkle::{int_to_leaf, int_to_node};
250
251 #[test]
252 fn get_root() {
253 let leaf0 = int_to_node(0);
254 let leaf1 = int_to_node(1);
255 let leaf2 = int_to_node(2);
256 let leaf3 = int_to_node(3);
257
258 let parent0 = calculate_parent_hash(leaf0, 0, leaf1);
259 let parent1 = calculate_parent_hash(leaf2, 2, leaf3);
260
261 let root_exp = calculate_parent_hash(parent0, 0, parent1);
262
263 let set = super::MerklePathSet::new(2)
264 .with_paths([(0, leaf0, vec![leaf1, parent1].into())])
265 .unwrap();
266
267 assert_eq!(set.root(), root_exp);
268 }
269
270 #[test]
271 fn add_and_get_path() {
272 let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)];
273 let hash_6 = int_to_node(6);
274 let index = 6_u64;
275 let depth = 3_u8;
276 let set = super::MerklePathSet::new(depth)
277 .with_paths([(index, hash_6, path_6.clone().into())])
278 .unwrap();
279 let stored_path_6 = set.get_path(NodeIndex::make(depth, index)).unwrap();
280
281 assert_eq!(path_6, *stored_path_6);
282 }
283
284 #[test]
285 fn get_node() {
286 let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)];
287 let hash_6 = int_to_node(6);
288 let index = 6_u64;
289 let depth = 3_u8;
290 let set = MerklePathSet::new(depth).with_paths([(index, hash_6, path_6.into())]).unwrap();
291
292 assert_eq!(int_to_node(6u64), set.get_node(NodeIndex::make(depth, index)).unwrap());
293 }
294
295 #[test]
296 fn update_leaf() {
297 let hash_4 = int_to_node(4);
298 let hash_5 = int_to_node(5);
299 let hash_6 = int_to_node(6);
300 let hash_7 = int_to_node(7);
301 let hash_45 = calculate_parent_hash(hash_4, 12u64, hash_5);
302 let hash_67 = calculate_parent_hash(hash_6, 14u64, hash_7);
303
304 let hash_0123 = int_to_node(123);
305
306 let path_6 = vec![hash_7, hash_45, hash_0123];
307 let path_5 = vec![hash_4, hash_67, hash_0123];
308 let path_4 = vec![hash_5, hash_67, hash_0123];
309
310 let index_6 = 6_u64;
311 let index_5 = 5_u64;
312 let index_4 = 4_u64;
313 let depth = 3_u8;
314 let mut set = MerklePathSet::new(depth)
315 .with_paths([
316 (index_6, hash_6, path_6.into()),
317 (index_5, hash_5, path_5.into()),
318 (index_4, hash_4, path_4.into()),
319 ])
320 .unwrap();
321
322 let new_hash_6 = int_to_leaf(100);
323 let new_hash_5 = int_to_leaf(55);
324
325 set.update_leaf(index_6, new_hash_6).unwrap();
326 let new_path_4 = set.get_path(NodeIndex::make(depth, index_4)).unwrap();
327 let new_hash_67 = calculate_parent_hash(new_hash_6.into(), 14_u64, hash_7);
328 assert_eq!(new_hash_67, new_path_4[1]);
329
330 set.update_leaf(index_5, new_hash_5).unwrap();
331 let new_path_4 = set.get_path(NodeIndex::make(depth, index_4)).unwrap();
332 let new_path_6 = set.get_path(NodeIndex::make(depth, index_6)).unwrap();
333 let new_hash_45 = calculate_parent_hash(new_hash_5.into(), 13_u64, hash_4);
334 assert_eq!(new_hash_45, new_path_6[1]);
335 assert_eq!(RpoDigest::from(new_hash_5), new_path_4[0]);
336 }
337
338 #[test]
339 fn depth_3_is_correct() {
340 let a = int_to_node(1);
341 let b = int_to_node(2);
342 let c = int_to_node(3);
343 let d = int_to_node(4);
344 let e = int_to_node(5);
345 let f = int_to_node(6);
346 let g = int_to_node(7);
347 let h = int_to_node(8);
348
349 let i = Rpo256::merge(&[a, b]);
350 let j = Rpo256::merge(&[c, d]);
351 let k = Rpo256::merge(&[e, f]);
352 let l = Rpo256::merge(&[g, h]);
353
354 let m = Rpo256::merge(&[i, j]);
355 let n = Rpo256::merge(&[k, l]);
356
357 let root = Rpo256::merge(&[m.into(), n.into()]);
358
359 let mut set = MerklePathSet::new(3);
360
361 let value = b;
362 let index = 1;
363 let path = MerklePath::new([a, j, n].to_vec());
364 set.add_path(index, value.into(), path.clone()).unwrap();
365 assert_eq!(*value, set.get_leaf(index).unwrap());
366 assert_eq!(root, set.root());
367
368 let value = e;
369 let index = 4;
370 let path = MerklePath::new([f.into(), l.into(), m.into()].to_vec());
371 set.add_path(index, value.into(), path.clone()).unwrap();
372 assert_eq!(*value, set.get_leaf(index).unwrap());
373 assert_eq!(root, set.root());
374
375 let value = a;
376 let index = 0;
377 let path = MerklePath::new([b.into(), j.into(), n.into()].to_vec());
378 set.add_path(index, value.into(), path.clone()).unwrap();
379 assert_eq!(*value, set.get_leaf(index).unwrap());
380 assert_eq!(root, set.root());
381
382 let value = h;
383 let index = 7;
384 let path = MerklePath::new([g.into(), k.into(), m.into()].to_vec());
385 set.add_path(index, value.into(), path.clone()).unwrap();
386 assert_eq!(*value, set.get_leaf(index).unwrap());
387 assert_eq!(root, set.root());
388 }
389
390 // HELPER FUNCTIONS
391 // --------------------------------------------------------------------------------------------
392
393 const fn is_even(pos: u64) -> bool {
394 pos & 1 == 0
395 }
396
397 /// Calculates the hash of the parent node by two sibling ones
398 /// - node — current node
399 /// - node_pos — position of the current node
400 /// - sibling — neighboring vertex in the tree
401 fn calculate_parent_hash(node: RpoDigest, node_pos: u64, sibling: RpoDigest) -> RpoDigest {
402 if is_even(node_pos) {
403 Rpo256::merge(&[node, sibling])
404 } else {
405 Rpo256::merge(&[sibling, node])
406 }
407 }
408}