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