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