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