miden-crypto/src/merkle/partial_mt/mod.rs
src/merkle/partial_mt/mod.rsRaw
1use super::{
2 BTreeMap, BTreeSet, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec,
3 Word, EMPTY_WORD,
4};
5
6#[cfg(test)]
7mod tests;
8
9// PARTIAL MERKLE TREE
10// ================================================================================================
11
12/// A partial Merkle tree with NodeIndex keys and 4-element RpoDigest leaf values.
13///
14/// The root of the tree is recomputed on each new leaf update.
15pub struct PartialMerkleTree {
16 max_depth: u8,
17 nodes: BTreeMap<NodeIndex, RpoDigest>,
18 leaves: BTreeSet<NodeIndex>,
19}
20
21impl Default for PartialMerkleTree {
22 fn default() -> Self {
23 Self::new()
24 }
25}
26
27impl PartialMerkleTree {
28 // CONSTANTS
29 // --------------------------------------------------------------------------------------------
30
31 /// An RpoDigest consisting of 4 ZERO elements.
32 pub const EMPTY_DIGEST: RpoDigest = RpoDigest::new(EMPTY_WORD);
33
34 /// Minimum supported depth.
35 pub const MIN_DEPTH: u8 = 1;
36
37 /// Maximum supported depth.
38 pub const MAX_DEPTH: u8 = 64;
39
40 pub const ROOT_INDEX: NodeIndex = NodeIndex::new_unchecked(0, 0);
41
42 // CONSTRUCTORS
43 // --------------------------------------------------------------------------------------------
44
45 /// Returns a new emply [PartialMerkleTree].
46 pub fn new() -> Self {
47 PartialMerkleTree {
48 max_depth: 0,
49 nodes: BTreeMap::new(),
50 leaves: BTreeSet::new(),
51 }
52 }
53
54 /// Appends the provided paths iterator into the set.
55 ///
56 /// Analogous to [Self::add_path].
57 pub fn with_paths<I>(paths: I) -> Result<Self, MerkleError>
58 where
59 I: IntoIterator<Item = (NodeIndex, Word, MerklePath)>,
60 {
61 // create an empty tree
62 let tree = PartialMerkleTree::new();
63
64 paths.into_iter().try_fold(tree, |mut tree, (index, value, path)| {
65 tree.add_path(index, value, path)?;
66 Ok(tree)
67 })
68 }
69
70 // PUBLIC ACCESSORS
71 // --------------------------------------------------------------------------------------------
72
73 /// Returns the root of this Merkle tree.
74 pub fn root(&self) -> Word {
75 *self.nodes.get(&Self::ROOT_INDEX).cloned().unwrap_or(Self::EMPTY_DIGEST)
76 }
77
78 /// Returns the depth of this Merkle tree.
79 pub fn max_depth(&self) -> u8 {
80 self.max_depth
81 }
82
83 /// Returns a node at the specified NodeIndex.
84 ///
85 /// # Errors
86 /// Returns an error if the specified NodeIndex is not contained in the nodes map.
87 pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
88 self.nodes.get(&index).ok_or(MerkleError::NodeNotInSet(index)).map(|hash| *hash)
89 }
90
91 /// Returns true if provided index contains in the leaves set, false otherwise.
92 pub fn is_leaf(&self, index: NodeIndex) -> bool {
93 self.leaves.contains(&index)
94 }
95
96 pub fn get_leaf_depth(&self, index: u64) -> Result<u8, MerkleError> {
97 let mut node_index = NodeIndex::new(self.max_depth(), index)?;
98 for _ in 0..node_index.depth() {
99 if self.leaves.contains(&node_index) {
100 return Ok(node_index.depth());
101 }
102 node_index.move_up()
103 }
104 // we don't have an error for this case, maybe it makes sense to create a new error, something like
105 // NoLeafForIndex("There is no leaf for provided index"). But it will be used almost never.
106 Err(MerkleError::NodeNotInSet(node_index))
107 }
108
109 /// Returns a value of the leaf at the specified NodeIndex.
110 ///
111 /// # Errors
112 /// Returns an error if the NodeIndex is not contained in the leaves set.
113 pub fn get_leaf(&self, index: NodeIndex) -> Result<Word, MerkleError> {
114 if !self.leaves.contains(&index) {
115 // This error not really suitable in this situation, should I create a new error?
116 Err(MerkleError::InvalidIndex {
117 depth: index.depth(),
118 value: index.value(),
119 })
120 } else {
121 self.nodes
122 .get(&index)
123 .ok_or(MerkleError::NodeNotInSet(index))
124 .map(|hash| **hash)
125 }
126 }
127
128 /// Returns a map of paths from every leaf to the root.
129 pub fn paths(&self) -> Result<BTreeMap<&NodeIndex, MerklePath>, MerkleError> {
130 let mut paths = BTreeMap::new();
131 for leaf_index in self.leaves.iter() {
132 let index = *leaf_index;
133 paths.insert(leaf_index, self.get_path(index)?);
134 }
135 Ok(paths)
136 }
137
138 /// Returns a Merkle path from the node at the specified index to the root.
139 ///
140 /// The node itself is not included in the path.
141 ///
142 /// # Errors
143 /// Returns an error if:
144 /// - the specified index has depth set to 0 or the depth is greater than the depth of this
145 /// Merkle tree.
146 /// - the specified index is not contained in the nodes map.
147 pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
148 if index.is_root() {
149 return Err(MerkleError::DepthTooSmall(index.depth()));
150 } else if index.depth() > self.max_depth() {
151 return Err(MerkleError::DepthTooBig(index.depth() as u64));
152 }
153
154 if !self.nodes.contains_key(&index) {
155 return Err(MerkleError::NodeNotInSet(index));
156 }
157
158 let mut path = Vec::new();
159 for _ in 0..index.depth() {
160 let sibling_index = Self::get_sibling_index(&index)?;
161 index.move_up();
162 let sibling_hash =
163 self.nodes.get(&sibling_index).cloned().unwrap_or(Self::EMPTY_DIGEST);
164 path.push(Word::from(sibling_hash));
165 }
166 Ok(MerklePath::new(path))
167 }
168
169 // ITERATORS
170 // --------------------------------------------------------------------------------------------
171
172 /// Returns an iterator over the leaves of this [PartialMerkleTree].
173 pub fn leaves(&self) -> impl Iterator<Item = (NodeIndex, &Word)> {
174 self.nodes
175 .iter()
176 .filter(|(index, _)| self.leaves.contains(index))
177 .map(|(index, hash)| (*index, &(**hash)))
178 }
179
180 /// Returns an iterator over the inner nodes of this Merkle tree.
181 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
182 let inner_nodes = self.nodes.iter().filter(|(index, _)| !self.leaves.contains(index));
183 inner_nodes.map(|(index, digest)| {
184 let left_index = NodeIndex::new(index.depth() + 1, index.value() * 2)
185 .expect("Failure to get left child index");
186 let right_index = NodeIndex::new(index.depth() + 1, index.value() * 2 + 1)
187 .expect("Failure to get right child index");
188 let left_hash = self.nodes.get(&left_index).cloned().unwrap_or(Self::EMPTY_DIGEST);
189 let right_hash = self.nodes.get(&right_index).cloned().unwrap_or(Self::EMPTY_DIGEST);
190 InnerNodeInfo {
191 value: **digest,
192 left: *left_hash,
193 right: *right_hash,
194 }
195 })
196 }
197
198 // STATE MUTATORS
199 // --------------------------------------------------------------------------------------------
200
201 /// Adds the nodes of the specified Merkle path to this [PartialMerkleTree]. The `index_value`
202 /// and `value` parameters specify the leaf node at which the path starts.
203 ///
204 /// # Errors
205 /// Returns an error if:
206 /// - The depth of the specified node_index is greater than 64 or smaller than 1.
207 /// - The specified path is not consistent with other paths in the set (i.e., resolves to a
208 /// different root).
209 pub fn add_path(
210 &mut self,
211 index_value: NodeIndex,
212 value: Word,
213 mut path: MerklePath,
214 ) -> Result<(), MerkleError> {
215 Self::check_depth(index_value.depth())?;
216 self.update_depth(index_value.depth());
217
218 // add node index to the leaves set
219 self.leaves.insert(index_value);
220 let sibling_node_index = Self::get_sibling_index(&index_value)?;
221 self.leaves.insert(sibling_node_index);
222
223 // add first two nodes to the nodes map
224 self.nodes.insert(index_value, value.into());
225 self.nodes.insert(sibling_node_index, path[0].into());
226
227 // update the current path
228 let parity = index_value.value() & 1;
229 path.insert(parity as usize, value);
230
231 // traverse to the root, updating the nodes
232 let mut index_value = index_value;
233 let root = Rpo256::merge(&[path[0].into(), path[1].into()]);
234 let root = path.iter().skip(2).copied().fold(root, |root, hash| {
235 index_value.move_up();
236 // insert calculated node to the nodes map
237 self.nodes.insert(index_value, root);
238
239 let sibling_node = Self::get_sibling_index_unchecked(&index_value);
240 // assume for now that all path nodes are leaves and add them to the leaves set
241 self.leaves.insert(sibling_node);
242
243 // insert node from Merkle path to the nodes map
244 self.nodes.insert(sibling_node, hash.into());
245
246 Rpo256::merge(&index_value.build_node(root, hash.into()))
247 });
248
249 let old_root = self.nodes.get(&Self::ROOT_INDEX).cloned().unwrap_or(Self::EMPTY_DIGEST);
250
251 // if the path set is empty (the root is all ZEROs), set the root to the root of the added
252 // path; otherwise, the root of the added path must be identical to the current root
253 if old_root == Self::EMPTY_DIGEST {
254 self.nodes.insert(Self::ROOT_INDEX, root);
255 } else if old_root != root {
256 return Err(MerkleError::ConflictingRoots([*old_root, *root].to_vec()));
257 }
258
259 self.update_leaves()?;
260
261 Ok(())
262 }
263
264 /// Updates value of the leaf at the specified index returning the old leaf value.
265 ///
266 /// This also recomputes all hashes between the leaf and the root, updating the root itself.
267 pub fn update_leaf(
268 &mut self,
269 node_index: NodeIndex,
270 value: RpoDigest,
271 ) -> Result<RpoDigest, MerkleError> {
272 // check correctness of the depth and update it
273 Self::check_depth(node_index.depth())?;
274 self.update_depth(node_index.depth());
275
276 // insert NodeIndex to the leaves Set
277 self.leaves.insert(node_index);
278
279 // add node value to the nodes Map
280 let old_value = self.nodes.insert(node_index, value).unwrap_or(Self::EMPTY_DIGEST);
281
282 // if the old value and new value are the same, there is nothing to update
283 if value == old_value {
284 return Ok(value);
285 }
286
287 let mut node_index = node_index;
288 let mut value = value;
289 for _ in 0..node_index.depth() {
290 let is_right = node_index.is_value_odd();
291 let (left, right) = if is_right {
292 let left_index = NodeIndex::new(node_index.depth(), node_index.value() - 1)?;
293 (
294 self.nodes
295 .get(&left_index)
296 .cloned()
297 .ok_or(MerkleError::NodeNotInSet(left_index))?,
298 value,
299 )
300 } else {
301 let right_index = NodeIndex::new(node_index.depth(), node_index.value() + 1)?;
302 (
303 value,
304 self.nodes
305 .get(&right_index)
306 .cloned()
307 .ok_or(MerkleError::NodeNotInSet(right_index))?,
308 )
309 };
310 node_index.move_up();
311 value = Rpo256::merge(&[left, right]);
312 self.nodes.insert(node_index, value);
313 }
314
315 Ok(old_value)
316 }
317
318 // HELPER METHODS
319 // --------------------------------------------------------------------------------------------
320
321 /// Updates depth value with the maximum of current and provided depth.
322 fn update_depth(&mut self, new_depth: u8) {
323 self.max_depth = new_depth.max(self.max_depth);
324 }
325
326 /// Returns an error if the depth is 0 or is greater than 64.
327 fn check_depth(depth: u8) -> Result<(), MerkleError> {
328 // validate the range of the depth.
329 if depth < Self::MIN_DEPTH {
330 return Err(MerkleError::DepthTooSmall(depth));
331 } else if Self::MAX_DEPTH < depth {
332 return Err(MerkleError::DepthTooBig(depth as u64));
333 }
334 Ok(())
335 }
336
337 fn get_sibling_index(node_index: &NodeIndex) -> Result<NodeIndex, MerkleError> {
338 if node_index.is_value_odd() {
339 NodeIndex::new(node_index.depth(), node_index.value() - 1)
340 } else {
341 NodeIndex::new(node_index.depth(), node_index.value() + 1)
342 }
343 }
344
345 fn get_sibling_index_unchecked(node_index: &NodeIndex) -> NodeIndex {
346 if node_index.is_value_odd() {
347 NodeIndex::new_unchecked(node_index.depth(), node_index.value() - 1)
348 } else {
349 NodeIndex::new_unchecked(node_index.depth(), node_index.value() + 1)
350 }
351 }
352
353 // Removes from the leaves set indexes of nodes which have descendants.
354 fn update_leaves(&mut self) -> Result<(), MerkleError> {
355 for leaf_node in self.leaves.clone().iter() {
356 let left_child = NodeIndex::new(leaf_node.depth() + 1, leaf_node.value() * 2)?;
357 let right_child = NodeIndex::new(leaf_node.depth() + 1, leaf_node.value() * 2 + 1)?;
358 if self.nodes.contains_key(&left_child) || self.nodes.contains_key(&right_child) {
359 self.leaves.remove(leaf_node);
360 }
361 }
362
363 Ok(())
364 }
365}