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