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