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