miden-crypto/src/merkle/path_set.rs
src/merkle/path_set.rsRaw
1use super::{BTreeMap, MerkleError, MerklePath, NodeIndex, 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: u8,
11 paths: BTreeMap<u64, MerklePath>,
12}
13
14impl MerklePathSet {
15 // CONSTRUCTOR
16 // --------------------------------------------------------------------------------------------
17
18 /// Returns an empty MerklePathSet.
19 pub fn new(depth: u8) -> 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) -> u8 {
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, index: NodeIndex) -> Result<Word, MerkleError> {
52 if !index.with_depth(self.total_depth).is_valid() {
53 return Err(MerkleError::InvalidIndex(
54 index.with_depth(self.total_depth),
55 ));
56 }
57 if index.depth() != self.total_depth {
58 return Err(MerkleError::InvalidDepth {
59 expected: self.total_depth,
60 provided: index.depth(),
61 });
62 }
63
64 let index_value = index.to_scalar_index();
65 let parity = index_value & 1;
66 let index_value = index_value / 2;
67 self.paths
68 .get(&index_value)
69 .ok_or(MerkleError::NodeNotInSet(index_value))
70 .map(|path| path[parity as usize])
71 }
72
73 /// Returns a Merkle path to the node at the specified index. The node itself is
74 /// not included in the path.
75 ///
76 /// # Errors
77 /// Returns an error if:
78 /// * The specified index not valid for the depth of structure.
79 /// * Node of the requested path does not exist in the set.
80 pub fn get_path(&self, index: NodeIndex) -> Result<MerklePath, MerkleError> {
81 if !index.with_depth(self.total_depth).is_valid() {
82 return Err(MerkleError::InvalidIndex(index));
83 }
84 if index.depth() != self.total_depth {
85 return Err(MerkleError::InvalidDepth {
86 expected: self.total_depth,
87 provided: index.depth(),
88 });
89 }
90
91 let index_value = index.to_scalar_index();
92 let index = index_value / 2;
93 let parity = index_value & 1;
94 let mut path = self
95 .paths
96 .get(&index)
97 .cloned()
98 .ok_or(MerkleError::NodeNotInSet(index))?;
99 path.remove(parity as usize);
100 Ok(path)
101 }
102
103 // STATE MUTATORS
104 // --------------------------------------------------------------------------------------------
105
106 /// Adds the specified Merkle path to this [MerklePathSet]. The `index` and `value` parameters
107 /// specify the leaf node at which the path starts.
108 ///
109 /// # Errors
110 /// Returns an error if:
111 /// - The specified index is not valid in the context of this Merkle path set (i.e., the index
112 /// implies a greater depth than is specified for this set).
113 /// - The specified path is not consistent with other paths in the set (i.e., resolves to a
114 /// different root).
115 pub fn add_path(
116 &mut self,
117 index_value: u64,
118 value: Word,
119 mut path: MerklePath,
120 ) -> Result<(), MerkleError> {
121 let depth = (path.len() + 1) as u8;
122 let mut index = NodeIndex::new(depth, index_value);
123 if index.depth() != self.total_depth {
124 return Err(MerkleError::InvalidDepth {
125 expected: self.total_depth,
126 provided: index.depth(),
127 });
128 }
129
130 // update the current path
131 let index_value = index.to_scalar_index();
132 let upper_index_value = index_value / 2;
133 let parity = index_value & 1;
134 path.insert(parity as usize, value);
135
136 // traverse to the root, updating the nodes
137 let root: Word = Rpo256::merge(&[path[0].into(), path[1].into()]).into();
138 let root = path.iter().skip(2).copied().fold(root, |root, hash| {
139 index.move_up();
140 Rpo256::merge(&index.build_node(root.into(), hash.into())).into()
141 });
142
143 // TODO review and document this logic
144 if self.root == [ZERO; 4] {
145 self.root = root;
146 } else if self.root != root {
147 return Err(MerkleError::InvalidPath(path));
148 }
149
150 // finish updating the path
151 self.paths.insert(upper_index_value, path);
152 Ok(())
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, base_index_value: u64, value: Word) -> Result<(), MerkleError> {
161 let depth = self.depth();
162 let mut index = NodeIndex::new(depth, base_index_value);
163 if !index.is_valid() {
164 return Err(MerkleError::InvalidIndex(index));
165 }
166
167 let path = match self
168 .paths
169 .get_mut(&index.clone().move_up().to_scalar_index())
170 {
171 Some(path) => path,
172 None => return Err(MerkleError::NodeNotInSet(base_index_value)),
173 };
174
175 // Fill old_hashes vector -----------------------------------------------------------------
176 let mut current_index = index;
177 let mut old_hashes = Vec::with_capacity(path.len().saturating_sub(2));
178 let mut root: Word = Rpo256::merge(&[path[0].into(), path[1].into()]).into();
179 for hash in path.iter().skip(2).copied() {
180 old_hashes.push(root);
181 current_index.move_up();
182 let input = current_index.build_node(hash.into(), root.into());
183 root = Rpo256::merge(&input).into();
184 }
185
186 // Fill new_hashes vector -----------------------------------------------------------------
187 path[index.is_value_odd() as usize] = value;
188
189 let mut new_hashes = Vec::with_capacity(path.len().saturating_sub(2));
190 let mut new_root: Word = Rpo256::merge(&[path[0].into(), path[1].into()]).into();
191 for path_hash in path.iter().skip(2).copied() {
192 new_hashes.push(new_root);
193 index.move_up();
194 let input = current_index.build_node(path_hash.into(), new_root.into());
195 new_root = Rpo256::merge(&input).into();
196 }
197
198 self.root = new_root;
199
200 // update paths ---------------------------------------------------------------------------
201 for path in self.paths.values_mut() {
202 for i in (0..old_hashes.len()).rev() {
203 if path[i + 2] == old_hashes[i] {
204 path[i + 2] = new_hashes[i];
205 break;
206 }
207 }
208 }
209
210 Ok(())
211 }
212}
213
214// TESTS
215// ================================================================================================
216
217#[cfg(test)]
218mod tests {
219 use super::*;
220 use crate::merkle::int_to_node;
221
222 #[test]
223 fn get_root() {
224 let leaf0 = int_to_node(0);
225 let leaf1 = int_to_node(1);
226 let leaf2 = int_to_node(2);
227 let leaf3 = int_to_node(3);
228
229 let parent0 = calculate_parent_hash(leaf0, 0, leaf1);
230 let parent1 = calculate_parent_hash(leaf2, 2, leaf3);
231
232 let root_exp = calculate_parent_hash(parent0, 0, parent1);
233
234 let mut set = super::MerklePathSet::new(3).unwrap();
235
236 set.add_path(0, leaf0, vec![leaf1, parent1].into()).unwrap();
237
238 assert_eq!(set.root(), root_exp);
239 }
240
241 #[test]
242 fn add_and_get_path() {
243 let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)];
244 let hash_6 = int_to_node(6);
245 let index = 6_u64;
246 let depth = 4_u8;
247 let mut set = super::MerklePathSet::new(depth).unwrap();
248
249 set.add_path(index, hash_6, path_6.clone().into()).unwrap();
250 let stored_path_6 = set.get_path(NodeIndex::new(depth, index)).unwrap();
251
252 assert_eq!(path_6, *stored_path_6);
253 assert!(set.get_path(NodeIndex::new(depth, 15_u64)).is_err())
254 }
255
256 #[test]
257 fn get_node() {
258 let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)];
259 let hash_6 = int_to_node(6);
260 let index = 6_u64;
261 let depth = 4_u8;
262 let mut set = MerklePathSet::new(depth).unwrap();
263
264 set.add_path(index, hash_6, path_6.into()).unwrap();
265
266 assert_eq!(
267 int_to_node(6u64),
268 set.get_node(NodeIndex::new(depth, index)).unwrap()
269 );
270 assert!(set.get_node(NodeIndex::new(depth, 15_u64)).is_err());
271 }
272
273 #[test]
274 fn update_leaf() {
275 let hash_4 = int_to_node(4);
276 let hash_5 = int_to_node(5);
277 let hash_6 = int_to_node(6);
278 let hash_7 = int_to_node(7);
279 let hash_45 = calculate_parent_hash(hash_4, 12u64, hash_5);
280 let hash_67 = calculate_parent_hash(hash_6, 14u64, hash_7);
281
282 let hash_0123 = int_to_node(123);
283
284 let path_6 = vec![hash_7, hash_45, hash_0123];
285 let path_5 = vec![hash_4, hash_67, hash_0123];
286 let path_4 = vec![hash_5, hash_67, hash_0123];
287
288 let index_6 = 6_u64;
289 let index_5 = 5_u64;
290 let index_4 = 4_u64;
291 let depth = 4_u8;
292 let mut set = MerklePathSet::new(depth).unwrap();
293
294 set.add_path(index_6, hash_6, path_6.into()).unwrap();
295 set.add_path(index_5, hash_5, path_5.into()).unwrap();
296 set.add_path(index_4, hash_4, path_4.into()).unwrap();
297
298 let new_hash_6 = int_to_node(100);
299 let new_hash_5 = int_to_node(55);
300
301 set.update_leaf(index_6, new_hash_6).unwrap();
302 let new_path_4 = set.get_path(NodeIndex::new(depth, index_4)).unwrap();
303 let new_hash_67 = calculate_parent_hash(new_hash_6, 14_u64, hash_7);
304 assert_eq!(new_hash_67, new_path_4[1]);
305
306 set.update_leaf(index_5, new_hash_5).unwrap();
307 let new_path_4 = set.get_path(NodeIndex::new(depth, index_4)).unwrap();
308 let new_path_6 = set.get_path(NodeIndex::new(depth, index_6)).unwrap();
309 let new_hash_45 = calculate_parent_hash(new_hash_5, 13_u64, hash_4);
310 assert_eq!(new_hash_45, new_path_6[1]);
311 assert_eq!(new_hash_5, new_path_4[0]);
312 }
313
314 // HELPER FUNCTIONS
315 // --------------------------------------------------------------------------------------------
316
317 const fn is_even(pos: u64) -> bool {
318 pos & 1 == 0
319 }
320
321 /// Calculates the hash of the parent node by two sibling ones
322 /// - node — current node
323 /// - node_pos — position of the current node
324 /// - sibling — neighboring vertex in the tree
325 fn calculate_parent_hash(node: Word, node_pos: u64, sibling: Word) -> Word {
326 if is_even(node_pos) {
327 Rpo256::merge(&[node.into(), sibling.into()]).into()
328 } else {
329 Rpo256::merge(&[sibling.into(), node.into()]).into()
330 }
331 }
332}