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