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