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).expect(&format!("Leaf with {leaf} is not in the nodes map")),
158 )
159 })
160 }
161
162 // STATE MUTATORS
163 // --------------------------------------------------------------------------------------------
164
165 /// Adds the nodes of the specified Merkle path to this [PartialMerkleTree]. The `index_value`
166 /// and `value` parameters specify the leaf node at which the path starts.
167 ///
168 /// # Errors
169 /// Returns an error if:
170 /// - The depth of the specified node_index is greater than 64 or smaller than 1.
171 /// - The specified path is not consistent with other paths in the set (i.e., resolves to a
172 /// different root).
173 pub fn add_path(
174 &mut self,
175 index_value: u64,
176 value: RpoDigest,
177 path: MerklePath,
178 ) -> Result<(), MerkleError> {
179 let index_value = NodeIndex::new(path.len() as u8, index_value)?;
180
181 Self::check_depth(index_value.depth())?;
182 self.update_depth(index_value.depth());
183
184 // add provided node and its sibling to the leaves set
185 self.leaves.insert(index_value);
186 let sibling_node_index = index_value.sibling();
187 self.leaves.insert(sibling_node_index);
188
189 // add provided node and its sibling to the nodes map
190 self.nodes.insert(index_value, value);
191 self.nodes.insert(sibling_node_index, path[0]);
192
193 // traverse to the root, updating the nodes
194 let mut index_value = index_value;
195 let node = Rpo256::merge(&index_value.build_node(value, path[0]));
196 let root = path.iter().skip(1).copied().fold(node, |node, hash| {
197 index_value.move_up();
198 // insert calculated node to the nodes map
199 self.nodes.insert(index_value, node);
200
201 // if the calculated node was a leaf, remove it from leaves set.
202 self.leaves.remove(&index_value);
203
204 let sibling_node = index_value.sibling();
205
206 // Insert node from Merkle path to the nodes map. This sibling node becomes a leaf only
207 // if it is a new node (it wasn't in nodes map).
208 // Node can be in 3 states: internal node, leaf of the tree and not a tree node at all.
209 // - Internal node can only stay in this state -- addition of a new path can't make it
210 // a leaf or remove it from the tree.
211 // - Leaf node can stay in the same state (remain a leaf) or can become an internal
212 // node. In the first case we don't need to do anything, and the second case is handled
213 // by the call of `self.leaves.remove(&index_value);`
214 // - New node can be a calculated node or a "sibling" node from a Merkle Path:
215 // --- Calculated node, obviously, never can be a leaf.
216 // --- Sibling node can be only a leaf, because otherwise it is not a new node.
217 if self.nodes.insert(sibling_node, hash).is_none() {
218 self.leaves.insert(sibling_node);
219 }
220
221 Rpo256::merge(&index_value.build_node(node, hash))
222 });
223
224 // if the path set is empty (the root is all ZEROs), set the root to the root of the added
225 // path; otherwise, the root of the added path must be identical to the current root
226 if self.root() == EMPTY_DIGEST {
227 self.nodes.insert(ROOT_INDEX, root);
228 } else if self.root() != root {
229 return Err(MerkleError::ConflictingRoots([self.root(), root].to_vec()));
230 }
231
232 Ok(())
233 }
234
235 /// Updates value of the leaf at the specified index returning the old leaf value.
236 ///
237 /// This also recomputes all hashes between the leaf and the root, updating the root itself.
238 ///
239 /// # Errors
240 /// Returns an error if:
241 /// - The depth of the specified node_index is greater than 64 or smaller than 1.
242 /// - The specified node index is not corresponding to the leaf.
243 pub fn update_leaf(
244 &mut self,
245 node_index: NodeIndex,
246 value: RpoDigest,
247 ) -> Result<RpoDigest, MerkleError> {
248 // check correctness of the depth and update it
249 Self::check_depth(node_index.depth())?;
250 self.update_depth(node_index.depth());
251
252 // insert NodeIndex to the leaves Set
253 self.leaves.insert(node_index);
254
255 // add node value to the nodes Map
256 let old_value = self
257 .nodes
258 .insert(node_index, value)
259 .ok_or(MerkleError::NodeNotInSet(node_index))?;
260
261 // if the old value and new value are the same, there is nothing to update
262 if value == old_value {
263 return Ok(old_value);
264 }
265
266 let mut node_index = node_index;
267 let mut value = value;
268 for _ in 0..node_index.depth() {
269 let sibling = self.nodes.get(&node_index.sibling()).expect("sibling should exist");
270 value = Rpo256::merge(&node_index.build_node(value, *sibling));
271 node_index.move_up();
272 self.nodes.insert(node_index, value);
273 }
274
275 Ok(old_value)
276 }
277
278 // UTILITY FUNCTIONS
279 // --------------------------------------------------------------------------------------------
280
281 /// Utility to visualize a [PartialMerkleTree] in text.
282 pub fn print(&self) -> Result<String, fmt::Error> {
283 let indent = " ";
284 let mut s = String::new();
285 s.push_str("root: ");
286 s.push_str(&word_to_hex(&self.root())?);
287 s.push('\n');
288 for d in 1..=self.max_depth() {
289 let entries = 2u64.pow(d.into());
290 for i in 0..entries {
291 let index = NodeIndex::new(d, i).expect("The index must always be valid");
292 let node = self.get_node(index);
293 let node = match node {
294 Err(_) => continue,
295 Ok(node) => node,
296 };
297
298 for _ in 0..d {
299 s.push_str(indent);
300 }
301 s.push_str(&format!("({}, {}): ", index.depth(), index.value()));
302 s.push_str(&word_to_hex(&node)?);
303 s.push('\n');
304 }
305 }
306
307 Ok(s)
308 }
309
310 // HELPER METHODS
311 // --------------------------------------------------------------------------------------------
312
313 /// Updates depth value with the maximum of current and provided depth.
314 fn update_depth(&mut self, new_depth: u8) {
315 self.max_depth = new_depth.max(self.max_depth);
316 }
317
318 /// Returns an error if the depth is 0 or is greater than 64.
319 fn check_depth(depth: u8) -> Result<(), MerkleError> {
320 // validate the range of the depth.
321 if depth < Self::MIN_DEPTH {
322 return Err(MerkleError::DepthTooSmall(depth));
323 } else if Self::MAX_DEPTH < depth {
324 return Err(MerkleError::DepthTooBig(depth as u64));
325 }
326 Ok(())
327 }
328}