miden-crypto/src/merkle/path_set.rs
src/merkle/path_set.rsRaw
1use super::{BTreeMap, MerkleError, MerklePath, NodeIndex, Rpo256, 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 all the leaf indexes of this path set.
61 pub fn indexes(&self) -> impl Iterator<Item = NodeIndex> + '_ {
62 self.paths
63 .keys()
64 .copied()
65 .map(|index| NodeIndex::new(self.total_depth, index))
66 }
67
68 /// Returns a node at the specified index.
69 ///
70 /// # Errors
71 /// Returns an error if:
72 /// * The specified index is not valid for the depth of structure.
73 /// * Requested node does not exist in the set.
74 pub fn get_node(&self, index: NodeIndex) -> Result<Word, MerkleError> {
75 if !index.with_depth(self.total_depth).is_valid() {
76 return Err(MerkleError::InvalidIndex(
77 index.with_depth(self.total_depth),
78 ));
79 }
80 if index.depth() != self.total_depth {
81 return Err(MerkleError::InvalidDepth {
82 expected: self.total_depth,
83 provided: index.depth(),
84 });
85 }
86
87 let index_value = index.to_scalar_index();
88 let parity = index_value & 1;
89 let index_value = index_value / 2;
90 self.paths
91 .get(&index_value)
92 .ok_or(MerkleError::NodeNotInSet(index_value))
93 .map(|path| path[parity as usize])
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 index_value = index.to_scalar_index();
115 let index = index_value / 2;
116 let parity = index_value & 1;
117 let mut path = self
118 .paths
119 .get(&index)
120 .cloned()
121 .ok_or(MerkleError::NodeNotInSet(index))?;
122 path.remove(parity as usize);
123 Ok(path)
124 }
125
126 // STATE MUTATORS
127 // --------------------------------------------------------------------------------------------
128
129 /// Adds the specified Merkle path to this [MerklePathSet]. The `index` and `value` parameters
130 /// specify the leaf node at which the path starts.
131 ///
132 /// # Errors
133 /// Returns an error if:
134 /// - The specified index is is not valid in the context of this Merkle path set (i.e., the
135 /// index implies a greater depth than is specified for this set).
136 /// - The specified path is not consistent with other paths in the set (i.e., resolves to a
137 /// different root).
138 pub fn add_path(
139 &mut self,
140 index_value: u64,
141 value: Word,
142 mut path: MerklePath,
143 ) -> Result<(), MerkleError> {
144 let depth = (path.len() + 1) as u8;
145 let mut index = NodeIndex::new(depth, index_value);
146 if index.depth() != self.total_depth {
147 return Err(MerkleError::InvalidDepth {
148 expected: self.total_depth,
149 provided: index.depth(),
150 });
151 }
152
153 // update the current path
154 let index_value = index.to_scalar_index();
155 let upper_index_value = index_value / 2;
156 let parity = index_value & 1;
157 path.insert(parity as usize, value);
158
159 // traverse to the root, updating the nodes
160 let root: Word = Rpo256::merge(&[path[0].into(), path[1].into()]).into();
161 let root = path.iter().skip(2).copied().fold(root, |root, hash| {
162 index.move_up();
163 Rpo256::merge(&index.build_node(root.into(), hash.into())).into()
164 });
165
166 // if the path set is empty (the root is all ZEROs), set the root to the root of the added
167 // path; otherwise, the root of the added path must be identical to the current root
168 if self.root == [ZERO; 4] {
169 self.root = root;
170 } else if self.root != root {
171 return Err(MerkleError::InvalidPath(path));
172 }
173
174 // finish updating the path
175 self.paths.insert(upper_index_value, path);
176 Ok(())
177 }
178
179 /// Replaces the leaf at the specified index with the provided value.
180 ///
181 /// # Errors
182 /// Returns an error if:
183 /// * Requested node does not exist in the set.
184 pub fn update_leaf(&mut self, base_index_value: u64, value: Word) -> Result<(), MerkleError> {
185 let depth = self.depth();
186 let mut index = NodeIndex::new(depth, base_index_value);
187 if !index.is_valid() {
188 return Err(MerkleError::InvalidIndex(index));
189 }
190
191 let path = match self
192 .paths
193 .get_mut(&index.clone().move_up().to_scalar_index())
194 {
195 Some(path) => path,
196 None => return Err(MerkleError::NodeNotInSet(base_index_value)),
197 };
198
199 // Fill old_hashes vector -----------------------------------------------------------------
200 let mut current_index = index;
201 let mut old_hashes = Vec::with_capacity(path.len().saturating_sub(2));
202 let mut root: Word = Rpo256::merge(&[path[0].into(), path[1].into()]).into();
203 for hash in path.iter().skip(2).copied() {
204 old_hashes.push(root);
205 current_index.move_up();
206 let input = current_index.build_node(hash.into(), root.into());
207 root = Rpo256::merge(&input).into();
208 }
209
210 // Fill new_hashes vector -----------------------------------------------------------------
211 path[index.is_value_odd() as usize] = value;
212
213 let mut new_hashes = Vec::with_capacity(path.len().saturating_sub(2));
214 let mut new_root: Word = Rpo256::merge(&[path[0].into(), path[1].into()]).into();
215 for path_hash in path.iter().skip(2).copied() {
216 new_hashes.push(new_root);
217 index.move_up();
218 let input = current_index.build_node(path_hash.into(), new_root.into());
219 new_root = Rpo256::merge(&input).into();
220 }
221
222 self.root = new_root;
223
224 // update paths ---------------------------------------------------------------------------
225 for path in self.paths.values_mut() {
226 for i in (0..old_hashes.len()).rev() {
227 if path[i + 2] == old_hashes[i] {
228 path[i + 2] = new_hashes[i];
229 break;
230 }
231 }
232 }
233
234 Ok(())
235 }
236}
237
238// TESTS
239// ================================================================================================
240
241#[cfg(test)]
242mod tests {
243 use super::*;
244 use crate::merkle::int_to_node;
245
246 #[test]
247 fn get_root() {
248 let leaf0 = int_to_node(0);
249 let leaf1 = int_to_node(1);
250 let leaf2 = int_to_node(2);
251 let leaf3 = int_to_node(3);
252
253 let parent0 = calculate_parent_hash(leaf0, 0, leaf1);
254 let parent1 = calculate_parent_hash(leaf2, 2, leaf3);
255
256 let root_exp = calculate_parent_hash(parent0, 0, parent1);
257
258 let set = super::MerklePathSet::new(3)
259 .with_paths([(0, leaf0, vec![leaf1, parent1].into())])
260 .unwrap();
261
262 assert_eq!(set.root(), root_exp);
263 }
264
265 #[test]
266 fn add_and_get_path() {
267 let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)];
268 let hash_6 = int_to_node(6);
269 let index = 6_u64;
270 let depth = 4_u8;
271 let set = super::MerklePathSet::new(depth)
272 .with_paths([(index, hash_6, path_6.clone().into())])
273 .unwrap();
274 let stored_path_6 = set.get_path(NodeIndex::new(depth, index)).unwrap();
275
276 assert_eq!(path_6, *stored_path_6);
277 assert!(set.get_path(NodeIndex::new(depth, 15_u64)).is_err())
278 }
279
280 #[test]
281 fn get_node() {
282 let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)];
283 let hash_6 = int_to_node(6);
284 let index = 6_u64;
285 let depth = 4_u8;
286 let set = MerklePathSet::new(depth)
287 .with_paths([(index, hash_6, path_6.into())])
288 .unwrap();
289
290 assert_eq!(
291 int_to_node(6u64),
292 set.get_node(NodeIndex::new(depth, index)).unwrap()
293 );
294 assert!(set.get_node(NodeIndex::new(depth, 15_u64)).is_err());
295 }
296
297 #[test]
298 fn update_leaf() {
299 let hash_4 = int_to_node(4);
300 let hash_5 = int_to_node(5);
301 let hash_6 = int_to_node(6);
302 let hash_7 = int_to_node(7);
303 let hash_45 = calculate_parent_hash(hash_4, 12u64, hash_5);
304 let hash_67 = calculate_parent_hash(hash_6, 14u64, hash_7);
305
306 let hash_0123 = int_to_node(123);
307
308 let path_6 = vec![hash_7, hash_45, hash_0123];
309 let path_5 = vec![hash_4, hash_67, hash_0123];
310 let path_4 = vec![hash_5, hash_67, hash_0123];
311
312 let index_6 = 6_u64;
313 let index_5 = 5_u64;
314 let index_4 = 4_u64;
315 let depth = 4_u8;
316 let mut set = MerklePathSet::new(depth)
317 .with_paths([
318 (index_6, hash_6, path_6.into()),
319 (index_5, hash_5, path_5.into()),
320 (index_4, hash_4, path_4.into()),
321 ])
322 .unwrap();
323
324 let new_hash_6 = int_to_node(100);
325 let new_hash_5 = int_to_node(55);
326
327 set.update_leaf(index_6, new_hash_6).unwrap();
328 let new_path_4 = set.get_path(NodeIndex::new(depth, index_4)).unwrap();
329 let new_hash_67 = calculate_parent_hash(new_hash_6, 14_u64, hash_7);
330 assert_eq!(new_hash_67, new_path_4[1]);
331
332 set.update_leaf(index_5, new_hash_5).unwrap();
333 let new_path_4 = set.get_path(NodeIndex::new(depth, index_4)).unwrap();
334 let new_path_6 = set.get_path(NodeIndex::new(depth, index_6)).unwrap();
335 let new_hash_45 = calculate_parent_hash(new_hash_5, 13_u64, hash_4);
336 assert_eq!(new_hash_45, new_path_6[1]);
337 assert_eq!(new_hash_5, new_path_4[0]);
338 }
339
340 // HELPER FUNCTIONS
341 // --------------------------------------------------------------------------------------------
342
343 const fn is_even(pos: u64) -> bool {
344 pos & 1 == 0
345 }
346
347 /// Calculates the hash of the parent node by two sibling ones
348 /// - node — current node
349 /// - node_pos — position of the current node
350 /// - sibling — neighboring vertex in the tree
351 fn calculate_parent_hash(node: Word, node_pos: u64, sibling: Word) -> Word {
352 if is_even(node_pos) {
353 Rpo256::merge(&[node.into(), sibling.into()]).into()
354 } else {
355 Rpo256::merge(&[sibling.into(), node.into()]).into()
356 }
357 }
358}