miden-crypto/src/merkle/tiered_smt/nodes.rs
src/merkle/tiered_smt/nodes.rsRaw
1use super::{
2 BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, LeafNodeIndex, MerkleError, MerklePath,
3 NodeIndex, Rpo256, RpoDigest, Vec,
4};
5
6// CONSTANTS
7// ================================================================================================
8
9/// The number of levels between tiers.
10const TIER_SIZE: u8 = super::TieredSmt::TIER_SIZE;
11
12/// Depths at which leaves can exist in a tiered SMT.
13const TIER_DEPTHS: [u8; 4] = super::TieredSmt::TIER_DEPTHS;
14
15/// Maximum node depth. This is also the bottom tier of the tree.
16const MAX_DEPTH: u8 = super::TieredSmt::MAX_DEPTH;
17
18// NODE STORE
19// ================================================================================================
20
21/// A store of nodes for a Tiered Sparse Merkle tree.
22///
23/// The store contains information about all nodes as well as information about which of the nodes
24/// represent leaf nodes in a Tiered Sparse Merkle tree. In the current implementation, [BTreeSet]s
25/// are used to determine the position of the leaves in the tree.
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub struct NodeStore {
28 nodes: BTreeMap<NodeIndex, RpoDigest>,
29 upper_leaves: BTreeSet<NodeIndex>,
30 bottom_leaves: BTreeSet<u64>,
31}
32
33impl NodeStore {
34 // CONSTRUCTOR
35 // --------------------------------------------------------------------------------------------
36 /// Returns a new instance of [NodeStore] instantiated with the specified root node.
37 ///
38 /// Root node is assumed to be a root of an empty sparse Merkle tree.
39 pub fn new(root_node: RpoDigest) -> Self {
40 let mut nodes = BTreeMap::default();
41 nodes.insert(NodeIndex::root(), root_node);
42
43 Self {
44 nodes,
45 upper_leaves: BTreeSet::default(),
46 bottom_leaves: BTreeSet::default(),
47 }
48 }
49
50 // PUBLIC ACCESSORS
51 // --------------------------------------------------------------------------------------------
52
53 /// Returns a node at the specified index.
54 ///
55 /// # Errors
56 /// Returns an error if:
57 /// - The specified index depth is 0 or greater than 64.
58 /// - The node with the specified index does not exists in the Merkle tree. This is possible
59 /// when a leaf node with the same index prefix exists at a tier higher than the requested
60 /// node.
61 pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
62 self.validate_node_access(index)?;
63 Ok(self.get_node_unchecked(&index))
64 }
65
66 /// Returns a Merkle path from the node at the specified index to the root.
67 ///
68 /// The node itself is not included in the path.
69 ///
70 /// # Errors
71 /// Returns an error if:
72 /// - The specified index depth is 0 or greater than 64.
73 /// - The node with the specified index does not exists in the Merkle tree. This is possible
74 /// when a leaf node with the same index prefix exists at a tier higher than the node to
75 /// which the path is requested.
76 pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
77 self.validate_node_access(index)?;
78
79 let mut path = Vec::with_capacity(index.depth() as usize);
80 for _ in 0..index.depth() {
81 let node = self.get_node_unchecked(&index.sibling());
82 path.push(node);
83 index.move_up();
84 }
85
86 Ok(path.into())
87 }
88
89 /// Returns a Merkle path to the node specified by the key together with a flag indicating,
90 /// whether this node is a leaf at depths 16, 32, or 48.
91 pub fn get_proof(&self, key: &RpoDigest) -> (MerklePath, NodeIndex, bool) {
92 let (index, leaf_exists) = self.get_leaf_index(key);
93 let index: NodeIndex = index.into();
94 let path = self.get_path(index).expect("failed to retrieve Merkle path for a node index");
95 (path, index, leaf_exists)
96 }
97
98 /// Returns an index at which a leaf node for the specified key should be inserted.
99 ///
100 /// The second value in the returned tuple is set to true if the node at the returned index
101 /// is already a leaf node.
102 pub fn get_leaf_index(&self, key: &RpoDigest) -> (LeafNodeIndex, bool) {
103 // traverse the tree from the root down checking nodes at tiers 16, 32, and 48. Return if
104 // a node at any of the tiers is either a leaf or a root of an empty subtree.
105 const NUM_UPPER_TIERS: usize = TIER_DEPTHS.len() - 1;
106 for &tier_depth in TIER_DEPTHS[..NUM_UPPER_TIERS].iter() {
107 let index = LeafNodeIndex::from_key(key, tier_depth);
108 if self.upper_leaves.contains(&index) {
109 return (index, true);
110 } else if !self.nodes.contains_key(&index) {
111 return (index, false);
112 }
113 }
114
115 // if we got here, that means all of the nodes checked so far are internal nodes, and
116 // the new node would need to be inserted in the bottom tier.
117 let index = LeafNodeIndex::from_key(key, MAX_DEPTH);
118 (index, self.bottom_leaves.contains(&index.value()))
119 }
120
121 /// Traverses the tree up from the bottom tier starting at the specified leaf index and
122 /// returns the depth of the first node which hash more than one child. The returned depth
123 /// is rounded up to the next tier.
124 pub fn get_last_single_child_parent_depth(&self, leaf_index: u64) -> u8 {
125 let mut index = NodeIndex::new_unchecked(MAX_DEPTH, leaf_index);
126
127 for _ in (TIER_DEPTHS[0]..MAX_DEPTH).rev() {
128 let sibling_index = index.sibling();
129 if self.nodes.contains_key(&sibling_index) {
130 break;
131 }
132 index.move_up();
133 }
134
135 let tier = (index.depth() - 1) / TIER_SIZE;
136 TIER_DEPTHS[tier as usize]
137 }
138
139 // ITERATORS
140 // --------------------------------------------------------------------------------------------
141
142 /// Returns an iterator over all inner nodes of the Tiered Sparse Merkle tree (i.e., nodes not
143 /// at depths 16 32, 48, or 64).
144 ///
145 /// The iterator order is unspecified.
146 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
147 self.nodes.iter().filter_map(|(index, node)| {
148 if self.is_internal_node(index) {
149 Some(InnerNodeInfo {
150 value: *node,
151 left: self.get_node_unchecked(&index.left_child()),
152 right: self.get_node_unchecked(&index.right_child()),
153 })
154 } else {
155 None
156 }
157 })
158 }
159
160 /// Returns an iterator over the upper leaves (i.e., leaves with depths 16, 32, 48) of the
161 /// Tiered Sparse Merkle tree.
162 pub fn upper_leaves(&self) -> impl Iterator<Item = (&NodeIndex, &RpoDigest)> {
163 self.upper_leaves.iter().map(|index| (index, &self.nodes[index]))
164 }
165
166 /// Returns an iterator over the bottom leaves (i.e., leaves with depth 64) of the Tiered
167 /// Sparse Merkle tree.
168 pub fn bottom_leaves(&self) -> impl Iterator<Item = (&u64, &RpoDigest)> {
169 self.bottom_leaves.iter().map(|value| {
170 let index = NodeIndex::new_unchecked(MAX_DEPTH, *value);
171 (value, &self.nodes[&index])
172 })
173 }
174
175 // STATE MUTATORS
176 // --------------------------------------------------------------------------------------------
177
178 /// Replaces the leaf node at the specified index with a tree consisting of two leaves located
179 /// at the specified indexes. Recomputes and returns the new root.
180 pub fn replace_leaf_with_subtree(
181 &mut self,
182 leaf_index: LeafNodeIndex,
183 subtree_leaves: [(LeafNodeIndex, RpoDigest); 2],
184 ) -> RpoDigest {
185 debug_assert!(self.is_non_empty_leaf(&leaf_index));
186 debug_assert!(!is_empty_root(&subtree_leaves[0].1));
187 debug_assert!(!is_empty_root(&subtree_leaves[1].1));
188 debug_assert_eq!(subtree_leaves[0].0.depth(), subtree_leaves[1].0.depth());
189 debug_assert!(leaf_index.depth() < subtree_leaves[0].0.depth());
190
191 self.upper_leaves.remove(&leaf_index);
192
193 if subtree_leaves[0].0 == subtree_leaves[1].0 {
194 // if the subtree is for a single node at depth 64, we only need to insert one node
195 debug_assert_eq!(subtree_leaves[0].0.depth(), MAX_DEPTH);
196 debug_assert_eq!(subtree_leaves[0].1, subtree_leaves[1].1);
197 self.insert_leaf_node(subtree_leaves[0].0, subtree_leaves[0].1)
198 } else {
199 self.insert_leaf_node(subtree_leaves[0].0, subtree_leaves[0].1);
200 self.insert_leaf_node(subtree_leaves[1].0, subtree_leaves[1].1)
201 }
202 }
203
204 /// Replaces a subtree containing the retained and the removed leaf nodes, with a leaf node
205 /// containing the retained leaf.
206 ///
207 /// This has the effect of deleting the the node at the `removed_leaf` index from the tree,
208 /// moving the node at the `retained_leaf` index up to the tier specified by `new_depth`.
209 pub fn replace_subtree_with_leaf(
210 &mut self,
211 removed_leaf: LeafNodeIndex,
212 retained_leaf: LeafNodeIndex,
213 new_depth: u8,
214 node: RpoDigest,
215 ) -> RpoDigest {
216 debug_assert!(!is_empty_root(&node));
217 debug_assert!(self.is_non_empty_leaf(&removed_leaf));
218 debug_assert!(self.is_non_empty_leaf(&retained_leaf));
219 debug_assert_eq!(removed_leaf.depth(), retained_leaf.depth());
220 debug_assert!(removed_leaf.depth() > new_depth);
221
222 // remove the branches leading up to the tier to which the retained leaf is to be moved
223 self.remove_branch(removed_leaf, new_depth);
224 self.remove_branch(retained_leaf, new_depth);
225
226 // compute the index of the common root for retained and removed leaves
227 let mut new_index = retained_leaf;
228 new_index.move_up_to(new_depth);
229
230 // insert the node at the root index
231 self.insert_leaf_node(new_index, node)
232 }
233
234 /// Inserts the specified node at the specified index; recomputes and returns the new root
235 /// of the Tiered Sparse Merkle tree.
236 ///
237 /// This method assumes that the provided node is a non-empty value, and that there is no node
238 /// at the specified index.
239 pub fn insert_leaf_node(&mut self, index: LeafNodeIndex, mut node: RpoDigest) -> RpoDigest {
240 debug_assert!(!is_empty_root(&node));
241 debug_assert_eq!(self.nodes.get(&index), None);
242
243 // mark the node as the leaf
244 if index.depth() == MAX_DEPTH {
245 self.bottom_leaves.insert(index.value());
246 } else {
247 self.upper_leaves.insert(index.into());
248 };
249
250 // insert the node and update the path from the node to the root
251 let mut index: NodeIndex = index.into();
252 for _ in 0..index.depth() {
253 self.nodes.insert(index, node);
254 let sibling = self.get_node_unchecked(&index.sibling());
255 node = Rpo256::merge(&index.build_node(node, sibling));
256 index.move_up();
257 }
258
259 // update the root
260 self.nodes.insert(NodeIndex::root(), node);
261 node
262 }
263
264 /// Updates the node at the specified index with the specified node value; recomputes and
265 /// returns the new root of the Tiered Sparse Merkle tree.
266 ///
267 /// This method can accept `node` as either an empty or a non-empty value.
268 pub fn update_leaf_node(&mut self, index: LeafNodeIndex, mut node: RpoDigest) -> RpoDigest {
269 debug_assert!(self.is_non_empty_leaf(&index));
270
271 // if the value we are updating the node to is a root of an empty tree, clear the leaf
272 // flag for this node
273 if node == EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize] {
274 if index.depth() == MAX_DEPTH {
275 self.bottom_leaves.remove(&index.value());
276 } else {
277 self.upper_leaves.remove(&index);
278 }
279 } else {
280 debug_assert!(!is_empty_root(&node));
281 }
282
283 // update the path from the node to the root
284 let mut index: NodeIndex = index.into();
285 for _ in 0..index.depth() {
286 if node == EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize] {
287 self.nodes.remove(&index);
288 } else {
289 self.nodes.insert(index, node);
290 }
291
292 let sibling = self.get_node_unchecked(&index.sibling());
293 node = Rpo256::merge(&index.build_node(node, sibling));
294 index.move_up();
295 }
296
297 // update the root
298 self.nodes.insert(NodeIndex::root(), node);
299 node
300 }
301
302 /// Replaces the leaf node at the specified index with a root of an empty subtree; recomputes
303 /// and returns the new root of the Tiered Sparse Merkle tree.
304 pub fn clear_leaf_node(&mut self, index: LeafNodeIndex) -> RpoDigest {
305 debug_assert!(self.is_non_empty_leaf(&index));
306 let node = EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize];
307 self.update_leaf_node(index, node)
308 }
309
310 /// Truncates a branch starting with specified leaf at the bottom tier to new depth.
311 ///
312 /// This involves removing the part of the branch below the new depth, and then inserting a new
313 /// // node at the new depth.
314 pub fn truncate_branch(
315 &mut self,
316 leaf_index: u64,
317 new_depth: u8,
318 node: RpoDigest,
319 ) -> RpoDigest {
320 debug_assert!(self.bottom_leaves.contains(&leaf_index));
321
322 let mut leaf_index = LeafNodeIndex::new(NodeIndex::new_unchecked(MAX_DEPTH, leaf_index));
323 self.remove_branch(leaf_index, new_depth);
324
325 leaf_index.move_up_to(new_depth);
326 self.insert_leaf_node(leaf_index, node)
327 }
328
329 // HELPER METHODS
330 // --------------------------------------------------------------------------------------------
331
332 /// Returns true if the node at the specified index is a leaf node.
333 fn is_non_empty_leaf(&self, index: &LeafNodeIndex) -> bool {
334 if index.depth() == MAX_DEPTH {
335 self.bottom_leaves.contains(&index.value())
336 } else {
337 self.upper_leaves.contains(index)
338 }
339 }
340
341 /// Returns true if the node at the specified index is an internal node - i.e., there is
342 /// no leaf at that node and the node does not belong to the bottom tier.
343 fn is_internal_node(&self, index: &NodeIndex) -> bool {
344 if index.depth() == MAX_DEPTH {
345 false
346 } else {
347 !self.upper_leaves.contains(index)
348 }
349 }
350
351 /// Checks if the specified index is valid in the context of this Merkle tree.
352 ///
353 /// # Errors
354 /// Returns an error if:
355 /// - The specified index depth is 0 or greater than 64.
356 /// - The node for the specified index does not exists in the Merkle tree. This is possible
357 /// when an ancestors of the specified index is a leaf node.
358 fn validate_node_access(&self, index: NodeIndex) -> Result<(), MerkleError> {
359 if index.is_root() {
360 return Err(MerkleError::DepthTooSmall(index.depth()));
361 } else if index.depth() > MAX_DEPTH {
362 return Err(MerkleError::DepthTooBig(index.depth() as u64));
363 } else {
364 // make sure that there are no leaf nodes in the ancestors of the index; since leaf
365 // nodes can live at specific depth, we just need to check these depths.
366 let tier = ((index.depth() - 1) / TIER_SIZE) as usize;
367 let mut tier_index = index;
368 for &depth in TIER_DEPTHS[..tier].iter().rev() {
369 tier_index.move_up_to(depth);
370 if self.upper_leaves.contains(&tier_index) {
371 return Err(MerkleError::NodeNotInSet(index));
372 }
373 }
374 }
375
376 Ok(())
377 }
378
379 /// Returns a node at the specified index. If the node does not exist at this index, a root
380 /// for an empty subtree at the index's depth is returned.
381 ///
382 /// Unlike [NodeStore::get_node()] this does not perform any checks to verify that the
383 /// returned node is valid in the context of this tree.
384 fn get_node_unchecked(&self, index: &NodeIndex) -> RpoDigest {
385 match self.nodes.get(index) {
386 Some(node) => *node,
387 None => EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize],
388 }
389 }
390
391 /// Removes a sequence of nodes starting at the specified index and traversing the tree up to
392 /// the specified depth. The node at the `end_depth` is also removed, and the appropriate leaf
393 /// flag is cleared.
394 ///
395 /// This method does not update any other nodes and does not recompute the tree root.
396 fn remove_branch(&mut self, index: LeafNodeIndex, end_depth: u8) {
397 if index.depth() == MAX_DEPTH {
398 self.bottom_leaves.remove(&index.value());
399 } else {
400 self.upper_leaves.remove(&index);
401 }
402
403 let mut index: NodeIndex = index.into();
404 assert!(index.depth() > end_depth);
405 for _ in 0..(index.depth() - end_depth + 1) {
406 self.nodes.remove(&index);
407 index.move_up()
408 }
409 }
410}
411
412// HELPER FUNCTIONS
413// ================================================================================================
414
415/// Returns true if the specified node is a root of an empty tree or an empty value ([ZERO; 4]).
416fn is_empty_root(node: &RpoDigest) -> bool {
417 EmptySubtreeRoots::empty_hashes(MAX_DEPTH).contains(node)
418}