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