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 // ITERATORS
122 // --------------------------------------------------------------------------------------------
123
124 /// Returns an iterator over all inner nodes of the Tiered Sparse Merkle tree (i.e., nodes not
125 /// at depths 16 32, 48, or 64).
126 ///
127 /// The iterator order is unspecified.
128 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
129 self.nodes.iter().filter_map(|(index, node)| {
130 if self.is_internal_node(index) {
131 Some(InnerNodeInfo {
132 value: *node,
133 left: self.get_node_unchecked(&index.left_child()),
134 right: self.get_node_unchecked(&index.right_child()),
135 })
136 } else {
137 None
138 }
139 })
140 }
141
142 /// Returns an iterator over the upper leaves (i.e., leaves with depths 16, 32, 48) of the
143 /// Tiered Sparse Merkle tree.
144 pub fn upper_leaves(&self) -> impl Iterator<Item = (&NodeIndex, &RpoDigest)> {
145 self.upper_leaves.iter().map(|index| (index, &self.nodes[index]))
146 }
147
148 /// Returns an iterator over the bottom leaves (i.e., leaves with depth 64) of the Tiered
149 /// Sparse Merkle tree.
150 pub fn bottom_leaves(&self) -> impl Iterator<Item = (&u64, &RpoDigest)> {
151 self.bottom_leaves.iter().map(|value| {
152 let index = NodeIndex::new_unchecked(MAX_DEPTH, *value);
153 (value, &self.nodes[&index])
154 })
155 }
156
157 // STATE MUTATORS
158 // --------------------------------------------------------------------------------------------
159
160 /// Replaces the leaf node at the specified index with a tree consisting of two leaves located
161 /// at the specified indexes. Recomputes and returns the new root.
162 pub fn replace_leaf_with_subtree(
163 &mut self,
164 leaf_index: LeafNodeIndex,
165 subtree_leaves: [(LeafNodeIndex, RpoDigest); 2],
166 ) -> RpoDigest {
167 debug_assert!(self.is_non_empty_leaf(&leaf_index));
168 debug_assert!(!is_empty_root(&subtree_leaves[0].1));
169 debug_assert!(!is_empty_root(&subtree_leaves[1].1));
170 debug_assert_eq!(subtree_leaves[0].0.depth(), subtree_leaves[1].0.depth());
171 debug_assert!(leaf_index.depth() < subtree_leaves[0].0.depth());
172
173 self.upper_leaves.remove(&leaf_index);
174
175 if subtree_leaves[0].0 == subtree_leaves[1].0 {
176 // if the subtree is for a single node at depth 64, we only need to insert one node
177 debug_assert_eq!(subtree_leaves[0].0.depth(), MAX_DEPTH);
178 debug_assert_eq!(subtree_leaves[0].1, subtree_leaves[1].1);
179 self.insert_leaf_node(subtree_leaves[0].0, subtree_leaves[0].1)
180 } else {
181 self.insert_leaf_node(subtree_leaves[0].0, subtree_leaves[0].1);
182 self.insert_leaf_node(subtree_leaves[1].0, subtree_leaves[1].1)
183 }
184 }
185
186 /// Replaces a subtree containing the retained and the removed leaf nodes, with a leaf node
187 /// containing the retained leaf.
188 ///
189 /// This has the effect of deleting the the node at the `removed_leaf` index from the tree,
190 /// moving the node at the `retained_leaf` index up to the tier specified by `new_depth`.
191 pub fn replace_subtree_with_leaf(
192 &mut self,
193 removed_leaf: LeafNodeIndex,
194 retained_leaf: LeafNodeIndex,
195 new_depth: u8,
196 node: RpoDigest,
197 ) -> RpoDigest {
198 debug_assert!(!is_empty_root(&node));
199 debug_assert!(self.is_non_empty_leaf(&removed_leaf));
200 debug_assert!(self.is_non_empty_leaf(&retained_leaf));
201 debug_assert_eq!(removed_leaf.depth(), retained_leaf.depth());
202 debug_assert!(removed_leaf.depth() > new_depth);
203
204 // clear leaf flags
205 if removed_leaf.depth() == MAX_DEPTH {
206 self.bottom_leaves.remove(&removed_leaf.value());
207 self.bottom_leaves.remove(&retained_leaf.value());
208 } else {
209 self.upper_leaves.remove(&removed_leaf);
210 self.upper_leaves.remove(&retained_leaf);
211 }
212
213 // remove the branches leading up to the tier to which the retained leaf is to be moved
214 self.remove_branch(removed_leaf, new_depth);
215 self.remove_branch(retained_leaf, new_depth);
216
217 // compute the index of the common root for retained and removed leaves
218 let mut new_index = retained_leaf;
219 new_index.move_up_to(new_depth);
220
221 // insert the node at the root index
222 self.insert_leaf_node(new_index, node)
223 }
224
225 /// Inserts the specified node at the specified index; recomputes and returns the new root
226 /// of the Tiered Sparse Merkle tree.
227 ///
228 /// This method assumes that the provided node is a non-empty value, and that there is no node
229 /// at the specified index.
230 pub fn insert_leaf_node(&mut self, index: LeafNodeIndex, mut node: RpoDigest) -> RpoDigest {
231 debug_assert!(!is_empty_root(&node));
232 debug_assert_eq!(self.nodes.get(&index), None);
233
234 // mark the node as the leaf
235 if index.depth() == MAX_DEPTH {
236 self.bottom_leaves.insert(index.value());
237 } else {
238 self.upper_leaves.insert(index.into());
239 };
240
241 // insert the node and update the path from the node to the root
242 let mut index: NodeIndex = index.into();
243 for _ in 0..index.depth() {
244 self.nodes.insert(index, node);
245 let sibling = self.get_node_unchecked(&index.sibling());
246 node = Rpo256::merge(&index.build_node(node, sibling));
247 index.move_up();
248 }
249
250 // update the root
251 self.nodes.insert(NodeIndex::root(), node);
252 node
253 }
254
255 /// Updates the node at the specified index with the specified node value; recomputes and
256 /// returns the new root of the Tiered Sparse Merkle tree.
257 ///
258 /// This method can accept `node` as either an empty or a non-empty value.
259 pub fn update_leaf_node(&mut self, index: LeafNodeIndex, mut node: RpoDigest) -> RpoDigest {
260 debug_assert!(self.is_non_empty_leaf(&index));
261
262 // if the value we are updating the node to is a root of an empty tree, clear the leaf
263 // flag for this node
264 if node == EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize] {
265 if index.depth() == MAX_DEPTH {
266 self.bottom_leaves.remove(&index.value());
267 } else {
268 self.upper_leaves.remove(&index);
269 }
270 } else {
271 debug_assert!(!is_empty_root(&node));
272 }
273
274 // update the path from the node to the root
275 let mut index: NodeIndex = index.into();
276 for _ in 0..index.depth() {
277 if node == EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize] {
278 self.nodes.remove(&index);
279 } else {
280 self.nodes.insert(index, node);
281 }
282
283 let sibling = self.get_node_unchecked(&index.sibling());
284 node = Rpo256::merge(&index.build_node(node, sibling));
285 index.move_up();
286 }
287
288 // update the root
289 self.nodes.insert(NodeIndex::root(), node);
290 node
291 }
292
293 /// Replaces the leaf node at the specified index with a root of an empty subtree; recomputes
294 /// and returns the new root of the Tiered Sparse Merkle tree.
295 pub fn clear_leaf_node(&mut self, index: LeafNodeIndex) -> RpoDigest {
296 debug_assert!(self.is_non_empty_leaf(&index));
297 let node = EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize];
298 self.update_leaf_node(index, node)
299 }
300
301 // HELPER METHODS
302 // --------------------------------------------------------------------------------------------
303
304 /// Returns true if the node at the specified index is a leaf node.
305 fn is_non_empty_leaf(&self, index: &LeafNodeIndex) -> bool {
306 if index.depth() == MAX_DEPTH {
307 self.bottom_leaves.contains(&index.value())
308 } else {
309 self.upper_leaves.contains(index)
310 }
311 }
312
313 /// Returns true if the node at the specified index is an internal node - i.e., there is
314 /// no leaf at that node and the node does not belong to the bottom tier.
315 fn is_internal_node(&self, index: &NodeIndex) -> bool {
316 if index.depth() == MAX_DEPTH {
317 false
318 } else {
319 !self.upper_leaves.contains(index)
320 }
321 }
322
323 /// Checks if the specified index is valid in the context of this Merkle tree.
324 ///
325 /// # Errors
326 /// Returns an error if:
327 /// - The specified index depth is 0 or greater than 64.
328 /// - The node for the specified index does not exists in the Merkle tree. This is possible
329 /// when an ancestors of the specified index is a leaf node.
330 fn validate_node_access(&self, index: NodeIndex) -> Result<(), MerkleError> {
331 if index.is_root() {
332 return Err(MerkleError::DepthTooSmall(index.depth()));
333 } else if index.depth() > MAX_DEPTH {
334 return Err(MerkleError::DepthTooBig(index.depth() as u64));
335 } else {
336 // make sure that there are no leaf nodes in the ancestors of the index; since leaf
337 // nodes can live at specific depth, we just need to check these depths.
338 let tier = ((index.depth() - 1) / TIER_SIZE) as usize;
339 let mut tier_index = index;
340 for &depth in TIER_DEPTHS[..tier].iter().rev() {
341 tier_index.move_up_to(depth);
342 if self.upper_leaves.contains(&tier_index) {
343 return Err(MerkleError::NodeNotInSet(index));
344 }
345 }
346 }
347
348 Ok(())
349 }
350
351 /// Returns a node at the specified index. If the node does not exist at this index, a root
352 /// for an empty subtree at the index's depth is returned.
353 ///
354 /// Unlike [NodeStore::get_node()] this does not perform any checks to verify that the
355 /// returned node is valid in the context of this tree.
356 fn get_node_unchecked(&self, index: &NodeIndex) -> RpoDigest {
357 match self.nodes.get(index) {
358 Some(node) => *node,
359 None => EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize],
360 }
361 }
362
363 /// Removes a sequence of nodes starting at the specified index and traversing the
364 /// tree up to the specified depth. The node at the `end_depth` is also removed.
365 ///
366 /// This method does not update any other nodes and does not recompute the tree root.
367 fn remove_branch(&mut self, index: LeafNodeIndex, end_depth: u8) {
368 let mut index: NodeIndex = index.into();
369 assert!(index.depth() > end_depth);
370 for _ in 0..(index.depth() - end_depth + 1) {
371 self.nodes.remove(&index);
372 index.move_up()
373 }
374 }
375}
376
377// HELPER FUNCTIONS
378// ================================================================================================
379
380/// Returns true if the specified node is a root of an empty tree or an empty value ([ZERO; 4]).
381fn is_empty_root(node: &RpoDigest) -> bool {
382 EmptySubtreeRoots::empty_hashes(MAX_DEPTH).contains(node)
383}