miden-crypto/src/merkle/path_set.rs
src/merkle/path_set.rsRaw
1use super::{BTreeMap, MerkleError, MerklePath, 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, MerklePath>,
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::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) -> u32 {
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 not valid for the depth of structure.
50 /// * Requested node does not exist in the set.
51 pub fn get_node(&self, depth: u32, index: u64) -> Result<Word, MerkleError> {
52 if index >= 2u64.pow(self.total_depth) {
53 return Err(MerkleError::InvalidIndex(self.total_depth, index));
54 }
55 if depth != self.total_depth {
56 return Err(MerkleError::InvalidDepth(self.total_depth, depth));
57 }
58
59 let pos = 2u64.pow(depth) + index;
60 let index = pos / 2;
61
62 match self.paths.get(&index) {
63 None => Err(MerkleError::NodeNotInSet(index)),
64 Some(path) => {
65 if Self::is_even(pos) {
66 Ok(path[0])
67 } else {
68 Ok(path[1])
69 }
70 }
71 }
72 }
73
74 /// Returns a Merkle path to the node at the specified index. The node itself is
75 /// not included in the path.
76 ///
77 /// # Errors
78 /// Returns an error if:
79 /// * The specified index not valid for the depth of structure.
80 /// * Node of the requested path does not exist in the set.
81 pub fn get_path(&self, depth: u32, index: u64) -> Result<MerklePath, MerkleError> {
82 if index >= 2u64.pow(self.total_depth) {
83 return Err(MerkleError::InvalidIndex(self.total_depth, index));
84 }
85 if depth != self.total_depth {
86 return Err(MerkleError::InvalidDepth(self.total_depth, depth));
87 }
88
89 let pos = 2u64.pow(depth) + index;
90 let index = pos / 2;
91
92 match self.paths.get(&index) {
93 None => Err(MerkleError::NodeNotInSet(index)),
94 Some(path) => {
95 let mut local_path = path.clone();
96 if Self::is_even(pos) {
97 local_path.remove(0);
98 Ok(local_path)
99 } else {
100 local_path.remove(1);
101 Ok(local_path)
102 }
103 }
104 }
105 }
106
107 // STATE MUTATORS
108 // --------------------------------------------------------------------------------------------
109
110 /// Adds the specified Merkle path to this [MerklePathSet]. The `index` and `value` parameters
111 /// specify the leaf node at which the path starts.
112 ///
113 /// # Errors
114 /// Returns an error if:
115 /// - The specified index is not valid in the context of this Merkle path set (i.e., the index
116 /// implies a greater depth than is specified for this set).
117 /// - The specified path is not consistent with other paths in the set (i.e., resolves to a
118 /// different root).
119 pub fn add_path(
120 &mut self,
121 index: u64,
122 value: Word,
123 path: MerklePath,
124 ) -> Result<(), MerkleError> {
125 let depth = (path.len() + 1) as u32;
126 if depth != self.total_depth {
127 return Err(MerkleError::InvalidDepth(self.total_depth, depth));
128 }
129
130 // Actual number of node in tree
131 let pos = 2u64.pow(self.total_depth) + index;
132
133 // Index of the leaf path in map. Paths of neighboring leaves are stored in one key-value pair
134 let half_pos = pos / 2;
135
136 let mut extended_path = path;
137 if Self::is_even(pos) {
138 extended_path.insert(0, value);
139 } else {
140 extended_path.insert(1, value);
141 }
142
143 let root_of_current_path = Self::compute_path_root(&extended_path, depth, index);
144 if self.root == [ZERO; 4] {
145 self.root = root_of_current_path;
146 } else if self.root != root_of_current_path {
147 return Err(MerkleError::InvalidPath(extended_path));
148 }
149 self.paths.insert(half_pos, extended_path);
150
151 Ok(())
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, _) = Self::compute_path_trace(path, depth, index);
173
174 // Fill new_hashes vector -----------------------------------------------------------------
175 if Self::is_even(pos) {
176 path[0] = value;
177 } else {
178 path[1] = value;
179 }
180
181 let (new_hashes, new_root) = Self::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 // HELPER FUNCTIONS
198 // --------------------------------------------------------------------------------------------
199
200 const fn is_even(pos: u64) -> bool {
201 pos & 1 == 0
202 }
203
204 /// Returns hash of the root
205 fn compute_path_root(path: &[Word], depth: u32, index: u64) -> Word {
206 let mut pos = 2u64.pow(depth) + index;
207
208 // hash that is obtained after calculating the current hash and path hash
209 let mut comp_hash = Rpo256::merge(&[path[0].into(), path[1].into()]).into();
210
211 for path_hash in path.iter().skip(2) {
212 pos /= 2;
213 comp_hash = Self::calculate_parent_hash(comp_hash, pos, *path_hash);
214 }
215
216 comp_hash
217 }
218
219 /// Calculates the hash of the parent node by two sibling ones
220 /// - node — current node
221 /// - node_pos — position of the current node
222 /// - sibling — neighboring vertex in the tree
223 fn calculate_parent_hash(node: Word, node_pos: u64, sibling: Word) -> Word {
224 if Self::is_even(node_pos) {
225 Rpo256::merge(&[node.into(), sibling.into()]).into()
226 } else {
227 Rpo256::merge(&[sibling.into(), node.into()]).into()
228 }
229 }
230
231 /// Returns vector of hashes from current to the root
232 fn compute_path_trace(path: &[Word], depth: u32, index: u64) -> (MerklePath, Word) {
233 let mut pos = 2u64.pow(depth) + index;
234
235 let mut computed_hashes = Vec::<Word>::new();
236
237 let mut comp_hash = Rpo256::merge(&[path[0].into(), path[1].into()]).into();
238
239 if path.len() != 2 {
240 for path_hash in path.iter().skip(2) {
241 computed_hashes.push(comp_hash);
242 pos /= 2;
243 comp_hash = Self::calculate_parent_hash(comp_hash, pos, *path_hash);
244 }
245 }
246
247 (computed_hashes.into(), comp_hash)
248 }
249}
250
251// TESTS
252// ================================================================================================
253
254#[cfg(test)]
255mod tests {
256 use super::*;
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 = MerklePathSet::calculate_parent_hash(leaf0, 0, leaf1);
267 let parent1 = MerklePathSet::calculate_parent_hash(leaf2, 2, leaf3);
268
269 let root_exp = MerklePathSet::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].into()).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().into()).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.into()).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 = MerklePathSet::calculate_parent_hash(hash_4, 12u64, hash_5);
314 let hash_67 = MerklePathSet::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.into()).unwrap();
329 set.add_path(index_5, hash_5, path_5.into()).unwrap();
330 set.add_path(index_4, hash_4, path_4.into()).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 = MerklePathSet::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 = MerklePathSet::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}