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