miden-crypto/crypto/src/merkle/merkle_path_set.rs
crypto/src/merkle/merkle_path_set.rsRaw
1use winter_utils::collections::BTreeMap;
2
3use super::{Felt, FieldElement, MerkleError, Word};
4use crate::hash::merge;
5
6// MERKLE PATH SET
7// ================================================================================================
8
9/// A set of Merkle paths.
10///
11/// This struct is intended to be used as one of the variants of the MerkleSet enum.
12#[derive(Clone, Debug)]
13pub struct MerklePathSet {
14 root: Word,
15 total_depth: u32,
16 paths: BTreeMap<u64, Vec<Word>>,
17}
18
19impl MerklePathSet {
20 // CONSTRUCTOR
21 // --------------------------------------------------------------------------------------------
22
23 /// Returns an empty MerklePathSet.
24 pub fn new(depth: u32) -> Result<Self, MerkleError> {
25 let root = [Felt::ZERO; 4];
26 let paths = BTreeMap::<u64, Vec<Word>>::new();
27
28 Ok(Self {
29 root,
30 total_depth: depth,
31 paths,
32 })
33 }
34
35 // PUBLIC ACCESSORS
36 // --------------------------------------------------------------------------------------------
37
38 /// Adds the specified Merkle path to this [MerklePathSet]. The `index` and `value` parameters
39 /// specify the leaf node at which the path starts.
40 ///
41 /// # Errors
42 /// Returns an error if:
43 /// - The specified index is not valid in the context of this Merkle path set (i.e., the index
44 /// implies a greater depth than is specified for this set).
45 /// - The specified path is not consistent with other paths in the set (i.e., resolves to a
46 /// different root).
47 pub fn add_path(
48 &mut self,
49 index: u64,
50 value: Word,
51 path: Vec<Word>,
52 ) -> Result<(), MerkleError> {
53 let depth = (path.len() + 1) as u32;
54 if depth != self.total_depth {
55 return Err(MerkleError::InvalidDepth(self.total_depth, depth));
56 }
57
58 // Actual number of node in tree
59 let pos = 2u64.pow(self.total_depth) + index;
60
61 // Index of the leaf path in map. Paths of neighboring leaves are stored in one key-value pair
62 let half_pos = (pos / 2) as u64;
63
64 let mut extended_path = path;
65 if is_even(pos) {
66 extended_path.insert(0, value);
67 } else {
68 extended_path.insert(1, value);
69 }
70
71 let root_of_current_path = compute_path_root(&extended_path, depth, index);
72 if self.root == [Felt::ZERO; 4] {
73 self.root = root_of_current_path;
74 } else if self.root != root_of_current_path {
75 return Err(MerkleError::InvalidPath(extended_path));
76 }
77 self.paths.insert(half_pos, extended_path);
78
79 Ok(())
80 }
81
82 /// Returns the root to which all paths in this set resolve.
83 pub fn root(&self) -> Word {
84 self.root
85 }
86
87 /// Returns the depth of the Merkle tree implied by the paths stored in this set.
88 ///
89 /// Merkle tree of depth 1 has two leaves, depth 2 has four leaves etc.
90 pub fn depth(&self) -> u32 {
91 self.total_depth
92 }
93
94 /// Returns a node at the specified index.
95 ///
96 /// # Errors
97 /// Returns an error if:
98 /// * The specified index not valid for the depth of structure.
99 /// * Requested node does not exist in the set.
100 pub fn get_node(&self, depth: u32, index: u64) -> Result<Word, MerkleError> {
101 if index >= 2u64.pow(self.total_depth) {
102 return Err(MerkleError::InvalidIndex(self.total_depth, index));
103 }
104 if depth != self.total_depth {
105 return Err(MerkleError::InvalidDepth(self.total_depth, depth));
106 }
107
108 let pos = 2u64.pow(depth) + index;
109 let index = (pos / 2) as u64;
110
111 match self.paths.get(&index) {
112 None => Err(MerkleError::NodeNotInSet(index)),
113 Some(path) => {
114 if is_even(pos) {
115 Ok(path[0])
116 } else {
117 Ok(path[1])
118 }
119 }
120 }
121 }
122
123 /// Returns a Merkle path to the node at the specified index. The node itself is
124 /// not included in the path.
125 ///
126 /// # Errors
127 /// Returns an error if:
128 /// * The specified index not valid for the depth of structure.
129 /// * Node of the requested path does not exist in the set.
130 pub fn get_path(&self, depth: u32, index: u64) -> Result<Vec<Word>, MerkleError> {
131 if index >= 2u64.pow(self.total_depth) {
132 return Err(MerkleError::InvalidIndex(self.total_depth, index));
133 }
134 if depth != self.total_depth {
135 return Err(MerkleError::InvalidDepth(self.total_depth, depth));
136 }
137
138 let pos = 2u64.pow(depth) + index;
139 let index = pos / 2;
140
141 match self.paths.get(&index) {
142 None => Err(MerkleError::NodeNotInSet(index)),
143 Some(path) => {
144 let mut local_path = path.clone();
145 if is_even(pos) {
146 local_path.remove(0);
147 Ok(local_path)
148 } else {
149 local_path.remove(1);
150 Ok(local_path)
151 }
152 }
153 }
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, index: u64, value: Word) -> Result<(), MerkleError> {
162 let depth = self.depth();
163 if index >= 2u64.pow(depth) {
164 return Err(MerkleError::InvalidIndex(depth, index));
165 }
166 let pos = 2u64.pow(depth) + index;
167
168 let path = match self.paths.get_mut(&(pos / 2)) {
169 None => return Err(MerkleError::NodeNotInSet(index)),
170 Some(path) => path,
171 };
172
173 // Fill old_hashes vector -----------------------------------------------------------------
174 let (old_hashes, _) = compute_path_trace(path, depth, index);
175
176 // Fill new_hashes vector -----------------------------------------------------------------
177 if is_even(pos) {
178 path[0] = value;
179 } else {
180 path[1] = value;
181 }
182
183 let (new_hashes, new_root) = compute_path_trace(path, depth, index);
184 self.root = new_root;
185
186 // update paths ---------------------------------------------------------------------------
187 for path in self.paths.values_mut() {
188 for i in (0..old_hashes.len()).rev() {
189 if path[i + 2] == old_hashes[i] {
190 path[i + 2] = new_hashes[i];
191 break;
192 }
193 }
194 }
195
196 Ok(())
197 }
198}
199
200// HELPER FUNCTIONS
201// --------------------------------------------------------------------------------------------
202
203fn is_even(pos: u64) -> bool {
204 pos & 1 == 0
205}
206
207/// Calculates the hash of the parent node by two sibling ones
208/// - node — current node
209/// - node_pos — position of the current node
210/// - sibling — neighboring vertex in the tree
211fn calculate_parent_hash(node: Word, node_pos: u64, sibling: Word) -> Word {
212 if is_even(node_pos) {
213 merge(&[node.into(), sibling.into()]).into()
214 } else {
215 merge(&[sibling.into(), node.into()]).into()
216 }
217}
218
219/// Returns vector of hashes from current to the root
220fn compute_path_trace(path: &[Word], depth: u32, index: u64) -> (Vec<Word>, Word) {
221 let mut pos = 2u64.pow(depth) + index;
222
223 let mut computed_hashes = Vec::<Word>::new();
224
225 let mut comp_hash = merge(&[path[0].into(), path[1].into()]).into();
226
227 if path.len() != 2 {
228 for path_hash in path.iter().skip(2) {
229 computed_hashes.push(comp_hash);
230 pos /= 2;
231 comp_hash = calculate_parent_hash(comp_hash, pos, *path_hash);
232 }
233 }
234
235 (computed_hashes, comp_hash)
236}
237
238/// Returns hash of the root
239fn compute_path_root(path: &[Word], depth: u32, index: u64) -> Word {
240 let mut pos = 2u64.pow(depth) + index;
241
242 // hash that is obtained after calculating the current hash and path hash
243 let mut comp_hash = merge(&[path[0].into(), path[1].into()]).into();
244
245 for path_hash in path.iter().skip(2) {
246 pos /= 2;
247 comp_hash = calculate_parent_hash(comp_hash, pos, *path_hash);
248 }
249
250 comp_hash
251}
252
253// TESTS
254// ================================================================================================
255
256#[cfg(test)]
257mod tests {
258 use super::calculate_parent_hash;
259
260 use super::{Felt, FieldElement, Word};
261
262 #[test]
263 fn get_root() {
264 let leaf0 = int_to_node(0);
265 let leaf1 = int_to_node(1);
266 let leaf2 = int_to_node(2);
267 let leaf3 = int_to_node(3);
268
269 let parent0 = calculate_parent_hash(leaf0, 0, leaf1);
270 let parent1 = calculate_parent_hash(leaf2, 2, leaf3);
271
272 let root_exp = calculate_parent_hash(parent0, 0, parent1);
273
274 let mut set = super::MerklePathSet::new(3).unwrap();
275
276 set.add_path(0, leaf0, vec![leaf1, parent1]).unwrap();
277
278 assert_eq!(set.root(), root_exp);
279 }
280
281 #[test]
282 fn add_and_get_path() {
283 let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)];
284 let hash_6 = int_to_node(6);
285 let index = 6u64;
286 let depth = 4u32;
287 let mut set = super::MerklePathSet::new(depth).unwrap();
288
289 set.add_path(index, hash_6, path_6.clone()).unwrap();
290 let stored_path_6 = set.get_path(depth, index).unwrap();
291
292 assert_eq!(path_6, stored_path_6);
293 assert!(set.get_path(depth, 15u64).is_err())
294 }
295
296 #[test]
297 fn get_node() {
298 let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)];
299 let hash_6 = int_to_node(6);
300 let index = 6u64;
301 let depth = 4u32;
302 let mut set = super::MerklePathSet::new(depth).unwrap();
303
304 set.add_path(index, hash_6, path_6).unwrap();
305
306 assert_eq!(int_to_node(6u64), set.get_node(depth, index).unwrap());
307 assert!(set.get_node(depth, 15u64).is_err());
308 }
309
310 #[test]
311 fn update_leaf() {
312 let hash_4 = int_to_node(4);
313 let hash_5 = int_to_node(5);
314 let hash_6 = int_to_node(6);
315 let hash_7 = int_to_node(7);
316 let hash_45 = calculate_parent_hash(hash_4, 12u64, hash_5);
317 let hash_67 = calculate_parent_hash(hash_6, 14u64, hash_7);
318
319 let hash_0123 = int_to_node(123);
320
321 let path_6 = vec![hash_7, hash_45, hash_0123];
322 let path_5 = vec![hash_4, hash_67, hash_0123];
323 let path_4 = vec![hash_5, hash_67, hash_0123];
324
325 let index_6 = 6u64;
326 let index_5 = 5u64;
327 let index_4 = 4u64;
328 let depth = 4u32;
329 let mut set = super::MerklePathSet::new(depth).unwrap();
330
331 set.add_path(index_6, hash_6, path_6).unwrap();
332 set.add_path(index_5, hash_5, path_5).unwrap();
333 set.add_path(index_4, hash_4, path_4).unwrap();
334
335 let new_hash_6 = int_to_node(100);
336 let new_hash_5 = int_to_node(55);
337
338 set.update_leaf(index_6, new_hash_6).unwrap();
339 let new_path_4 = set.get_path(depth, index_4).unwrap();
340 let new_hash_67 = calculate_parent_hash(new_hash_6, 14u64, hash_7);
341 assert_eq!(new_hash_67, new_path_4[1]);
342
343 set.update_leaf(index_5, new_hash_5).unwrap();
344 let new_path_4 = set.get_path(depth, index_4).unwrap();
345 let new_path_6 = set.get_path(depth, index_6).unwrap();
346 let new_hash_45 = calculate_parent_hash(new_hash_5, 13u64, hash_4);
347 assert_eq!(new_hash_45, new_path_6[1]);
348 assert_eq!(new_hash_5, new_path_4[0]);
349 }
350
351 // HELPER FUNCTIONS
352 // --------------------------------------------------------------------------------------------
353
354 const fn int_to_node(value: u64) -> Word {
355 [Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO]
356 }
357}