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