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