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