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