miden-crypto/src/merkle/tiered_smt/nodes.rs
src/merkle/tiered_smt/nodes.rsRaw
1use super::{
2 get_index_tier, get_key_prefix, is_leaf_node, BTreeMap, BTreeSet, EmptySubtreeRoots,
3 InnerNodeInfo, MerkleError, MerklePath, 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.
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct NodeStore {
27 nodes: BTreeMap<NodeIndex, RpoDigest>,
28 upper_leaves: BTreeSet<NodeIndex>,
29 bottom_leaves: BTreeSet<u64>,
30}
31
32impl NodeStore {
33 // CONSTRUCTOR
34 // --------------------------------------------------------------------------------------------
35 /// Returns a new instance of [NodeStore] instantiated with the specified root node.
36 ///
37 /// Root node is assumed to be a root of an empty sparse Merkle tree.
38 pub fn new(root_node: RpoDigest) -> Self {
39 let mut nodes = BTreeMap::default();
40 nodes.insert(NodeIndex::root(), root_node);
41
42 Self {
43 nodes,
44 upper_leaves: BTreeSet::default(),
45 bottom_leaves: BTreeSet::default(),
46 }
47 }
48
49 // PUBLIC ACCESSORS
50 // --------------------------------------------------------------------------------------------
51
52 /// Returns a node at the specified index.
53 ///
54 /// # Errors
55 /// Returns an error if:
56 /// - The specified index depth is 0 or greater than 64.
57 /// - The node with the specified index does not exists in the Merkle tree. This is possible
58 /// when a leaf node with the same index prefix exists at a tier higher than the requested
59 /// node.
60 pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
61 self.validate_node_access(index)?;
62 Ok(self.get_node_unchecked(&index))
63 }
64
65 /// Returns a Merkle path from the node at the specified index to the root.
66 ///
67 /// The node itself is not included in the path.
68 ///
69 /// # Errors
70 /// Returns an error if:
71 /// - The specified index depth is 0 or greater than 64.
72 /// - The node with the specified index does not exists in the Merkle tree. This is possible
73 /// when a leaf node with the same index prefix exists at a tier higher than the node to
74 /// which the path is requested.
75 pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
76 self.validate_node_access(index)?;
77
78 let mut path = Vec::with_capacity(index.depth() as usize);
79 for _ in 0..index.depth() {
80 let node = self.get_node_unchecked(&index.sibling());
81 path.push(node);
82 index.move_up();
83 }
84
85 Ok(path.into())
86 }
87
88 /// Returns an index at which a leaf node for the specified key should be inserted.
89 ///
90 /// The second value in the returned tuple is set to true if the node at the returned index
91 /// is already a leaf node, excluding leaves at the bottom tier (i.e., if the leaf is at the
92 /// bottom tier, false is returned).
93 pub fn get_insert_location(&self, key: &RpoDigest) -> (NodeIndex, 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 let mse = get_key_prefix(key);
97 for depth in (TIER_DEPTHS[0]..MAX_DEPTH).step_by(TIER_SIZE as usize) {
98 let index = NodeIndex::new_unchecked(depth, mse >> (MAX_DEPTH - 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 = NodeIndex::new_unchecked(MAX_DEPTH, mse);
109 (index, false)
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 !is_leaf_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: NodeIndex,
156 subtree_leaves: [(NodeIndex, RpoDigest); 2],
157 ) -> RpoDigest {
158 debug_assert!(is_leaf_node(&leaf_index));
159 debug_assert!(is_leaf_node(&subtree_leaves[0].0));
160 debug_assert!(is_leaf_node(&subtree_leaves[1].0));
161 debug_assert!(!is_empty_root(&subtree_leaves[0].1));
162 debug_assert!(!is_empty_root(&subtree_leaves[1].1));
163 debug_assert_eq!(subtree_leaves[0].0.depth(), subtree_leaves[1].0.depth());
164 debug_assert!(leaf_index.depth() < subtree_leaves[0].0.depth());
165
166 self.upper_leaves.remove(&leaf_index);
167 self.insert_leaf_node(subtree_leaves[0].0, subtree_leaves[0].1);
168 self.insert_leaf_node(subtree_leaves[1].0, subtree_leaves[1].1)
169 }
170
171 /// Replaces a subtree containing the retained and the removed leaf nodes, with a leaf node
172 /// containing the retained leaf.
173 ///
174 /// This has the effect of deleting the the node at the `removed_leaf` index from the tree,
175 /// moving the node at the `retained_leaf` index up to the tier specified by `new_depth`.
176 pub fn replace_subtree_with_leaf(
177 &mut self,
178 removed_leaf: NodeIndex,
179 retained_leaf: NodeIndex,
180 new_depth: u8,
181 node: RpoDigest,
182 ) -> RpoDigest {
183 debug_assert!(!is_empty_root(&node));
184 debug_assert!(self.is_leaf(&removed_leaf));
185 debug_assert!(self.is_leaf(&retained_leaf));
186 debug_assert_eq!(removed_leaf.depth(), retained_leaf.depth());
187 debug_assert!(removed_leaf.depth() > new_depth);
188
189 // clear leaf flags
190 if removed_leaf.depth() == MAX_DEPTH {
191 self.bottom_leaves.remove(&removed_leaf.value());
192 self.bottom_leaves.remove(&retained_leaf.value());
193 } else {
194 self.upper_leaves.remove(&removed_leaf);
195 self.upper_leaves.remove(&retained_leaf);
196 }
197
198 // remove the branches leading up to the tier to which the retained leaf is to be moved
199 self.remove_branch(removed_leaf, new_depth);
200 self.remove_branch(retained_leaf, new_depth);
201
202 // compute the index of the common root for retained and removed leaves
203 let mut new_index = retained_leaf;
204 new_index.move_up_to(new_depth);
205 debug_assert!(is_leaf_node(&new_index));
206
207 // insert the node at the root index
208 self.insert_leaf_node(new_index, node)
209 }
210
211 /// Inserts the specified node at the specified index; recomputes and returns the new root
212 /// of the Tiered Sparse Merkle tree.
213 ///
214 /// This method assumes that node is a non-empty value.
215 pub fn insert_leaf_node(&mut self, mut index: NodeIndex, mut node: RpoDigest) -> RpoDigest {
216 debug_assert!(is_leaf_node(&index));
217 debug_assert!(!is_empty_root(&node));
218
219 // mark the node as the leaf
220 if index.depth() == MAX_DEPTH {
221 self.bottom_leaves.insert(index.value());
222 } else {
223 self.upper_leaves.insert(index);
224 };
225
226 // insert the node and update the path from the node to the root
227 for _ in 0..index.depth() {
228 self.nodes.insert(index, node);
229 let sibling = self.get_node_unchecked(&index.sibling());
230 node = Rpo256::merge(&index.build_node(node, sibling));
231 index.move_up();
232 }
233
234 // update the root
235 self.nodes.insert(NodeIndex::root(), node);
236 node
237 }
238
239 /// Updates the node at the specified index with the specified node value; recomputes and
240 /// returns the new root of the Tiered Sparse Merkle tree.
241 ///
242 /// This method can accept `node` as either an empty or a non-empty value.
243 pub fn update_leaf_node(&mut self, mut index: NodeIndex, mut node: RpoDigest) -> RpoDigest {
244 debug_assert!(self.is_leaf(&index));
245
246 // if the value we are updating the node to is a root of an empty tree, clear the leaf
247 // flag for this node
248 if node == EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize] {
249 if index.depth() == MAX_DEPTH {
250 self.bottom_leaves.remove(&index.value());
251 } else {
252 self.upper_leaves.remove(&index);
253 }
254 } else {
255 debug_assert!(!is_empty_root(&node));
256 }
257
258 // update the path from the node to the root
259 for _ in 0..index.depth() {
260 if node == EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize] {
261 self.nodes.remove(&index);
262 } else {
263 self.nodes.insert(index, node);
264 }
265
266 let sibling = self.get_node_unchecked(&index.sibling());
267 node = Rpo256::merge(&index.build_node(node, sibling));
268 index.move_up();
269 }
270
271 // update the root
272 self.nodes.insert(NodeIndex::root(), node);
273 node
274 }
275
276 /// Replaces the leaf node at the specified index with a root of an empty subtree; recomputes
277 /// and returns the new root of the Tiered Sparse Merkle tree.
278 pub fn clear_leaf_node(&mut self, index: NodeIndex) -> RpoDigest {
279 debug_assert!(self.is_leaf(&index));
280 let node = EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize];
281 self.update_leaf_node(index, node)
282 }
283
284 // HELPER METHODS
285 // --------------------------------------------------------------------------------------------
286
287 /// Returns true if the node at the specified index is a leaf node.
288 fn is_leaf(&self, index: &NodeIndex) -> bool {
289 debug_assert!(is_leaf_node(index));
290 if index.depth() == MAX_DEPTH {
291 self.bottom_leaves.contains(&index.value())
292 } else {
293 self.upper_leaves.contains(index)
294 }
295 }
296
297 /// Checks if the specified index is valid in the context of this Merkle tree.
298 ///
299 /// # Errors
300 /// Returns an error if:
301 /// - The specified index depth is 0 or greater than 64.
302 /// - The node for the specified index does not exists in the Merkle tree. This is possible
303 /// when an ancestors of the specified index is a leaf node.
304 fn validate_node_access(&self, index: NodeIndex) -> Result<(), MerkleError> {
305 if index.is_root() {
306 return Err(MerkleError::DepthTooSmall(index.depth()));
307 } else if index.depth() > MAX_DEPTH {
308 return Err(MerkleError::DepthTooBig(index.depth() as u64));
309 } else {
310 // make sure that there are no leaf nodes in the ancestors of the index; since leaf
311 // nodes can live at specific depth, we just need to check these depths.
312 let tier = get_index_tier(&index);
313 let mut tier_index = index;
314 for &depth in TIER_DEPTHS[..tier].iter().rev() {
315 tier_index.move_up_to(depth);
316 if self.upper_leaves.contains(&tier_index) {
317 return Err(MerkleError::NodeNotInSet(index));
318 }
319 }
320 }
321
322 Ok(())
323 }
324
325 /// Returns a node at the specified index. If the node does not exist at this index, a root
326 /// for an empty subtree at the index's depth is returned.
327 ///
328 /// Unlike [NodeStore::get_node()] this does not perform any checks to verify that the
329 /// returned node is valid in the context of this tree.
330 fn get_node_unchecked(&self, index: &NodeIndex) -> RpoDigest {
331 match self.nodes.get(index) {
332 Some(node) => *node,
333 None => EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize],
334 }
335 }
336
337 /// Removes a sequence of nodes starting at the specified index and traversing the
338 /// tree up to the specified depth.
339 ///
340 /// This method does not update any other nodes and does not recompute the tree root.
341 fn remove_branch(&mut self, mut index: NodeIndex, end_depth: u8) {
342 assert!(index.depth() > end_depth);
343 for _ in 0..(index.depth() - end_depth) {
344 self.nodes.remove(&index);
345 index.move_up()
346 }
347 }
348}
349
350// HELPER FUNCTIONS
351// ================================================================================================
352
353/// Returns true if the specified node is a root of an empty tree or an empty value ([ZERO; 4]).
354fn is_empty_root(node: &RpoDigest) -> bool {
355 EmptySubtreeRoots::empty_hashes(MAX_DEPTH).contains(node)
356}