miden-crypto/src/merkle/simple_smt/mod.rs
src/merkle/simple_smt/mod.rsRaw
1use super::{
2 BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerkleTreeDelta,
3 NodeIndex, Rpo256, RpoDigest, StoreNode, TryApplyDiff, Vec, Word,
4};
5
6#[cfg(test)]
7mod tests;
8
9// SPARSE MERKLE TREE
10// ================================================================================================
11
12/// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction.
13///
14/// The root of the tree is recomputed on each new leaf update.
15#[derive(Debug, Clone, PartialEq, Eq)]
16#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
17pub struct SimpleSmt {
18 depth: u8,
19 root: RpoDigest,
20 leaves: BTreeMap<u64, Word>,
21 branches: BTreeMap<NodeIndex, BranchNode>,
22}
23
24impl SimpleSmt {
25 // CONSTANTS
26 // --------------------------------------------------------------------------------------------
27
28 /// Minimum supported depth.
29 pub const MIN_DEPTH: u8 = 1;
30
31 /// Maximum supported depth.
32 pub const MAX_DEPTH: u8 = 64;
33
34 /// Value of an empty leaf.
35 pub const EMPTY_VALUE: Word = super::EMPTY_WORD;
36
37 // CONSTRUCTORS
38 // --------------------------------------------------------------------------------------------
39
40 /// Returns a new [SimpleSmt] instantiated with the specified depth.
41 ///
42 /// All leaves in the returned tree are set to [ZERO; 4].
43 ///
44 /// # Errors
45 /// Returns an error if the depth is 0 or is greater than 64.
46 pub fn new(depth: u8) -> Result<Self, MerkleError> {
47 // validate the range of the depth.
48 if depth < Self::MIN_DEPTH {
49 return Err(MerkleError::DepthTooSmall(depth));
50 } else if Self::MAX_DEPTH < depth {
51 return Err(MerkleError::DepthTooBig(depth as u64));
52 }
53
54 let root = *EmptySubtreeRoots::entry(depth, 0);
55
56 Ok(Self {
57 root,
58 depth,
59 leaves: BTreeMap::new(),
60 branches: BTreeMap::new(),
61 })
62 }
63
64 /// Returns a new [SimpleSmt] instantiated with the specified depth and with leaves
65 /// set as specified by the provided entries.
66 ///
67 /// All leaves omitted from the entries list are set to [ZERO; 4].
68 ///
69 /// # Errors
70 /// Returns an error if:
71 /// - If the depth is 0 or is greater than 64.
72 /// - The number of entries exceeds the maximum tree capacity, that is 2^{depth}.
73 /// - The provided entries contain multiple values for the same key.
74 pub fn with_leaves(
75 depth: u8,
76 entries: impl IntoIterator<Item = (u64, Word)>,
77 ) -> Result<Self, MerkleError> {
78 // create an empty tree
79 let mut tree = Self::new(depth)?;
80
81 // compute the max number of entries. We use an upper bound of depth 63 because we consider
82 // passing in a vector of size 2^64 infeasible.
83 let max_num_entries = 2_usize.pow(tree.depth.min(63).into());
84
85 // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
86 // entries with the empty value need additional tracking.
87 let mut key_set_to_zero = BTreeSet::new();
88
89 for (idx, (key, value)) in entries.into_iter().enumerate() {
90 if idx >= max_num_entries {
91 return Err(MerkleError::InvalidNumEntries(max_num_entries));
92 }
93
94 let old_value = tree.update_leaf(key, value)?;
95
96 if old_value != Self::EMPTY_VALUE || key_set_to_zero.contains(&key) {
97 return Err(MerkleError::DuplicateValuesForIndex(key));
98 }
99
100 if value == Self::EMPTY_VALUE {
101 key_set_to_zero.insert(key);
102 };
103 }
104 Ok(tree)
105 }
106
107 /// Wrapper around [`SimpleSmt::with_leaves`] which inserts leaves at contiguous indices
108 /// starting at index 0.
109 pub fn with_contiguous_leaves(
110 depth: u8,
111 entries: impl IntoIterator<Item = Word>,
112 ) -> Result<Self, MerkleError> {
113 Self::with_leaves(
114 depth,
115 entries
116 .into_iter()
117 .enumerate()
118 .map(|(idx, word)| (idx.try_into().expect("tree max depth is 2^8"), word)),
119 )
120 }
121
122 // PUBLIC ACCESSORS
123 // --------------------------------------------------------------------------------------------
124
125 /// Returns the root of this Merkle tree.
126 pub const fn root(&self) -> RpoDigest {
127 self.root
128 }
129
130 /// Returns the depth of this Merkle tree.
131 pub const fn depth(&self) -> u8 {
132 self.depth
133 }
134
135 /// Returns a node at the specified index.
136 ///
137 /// # Errors
138 /// Returns an error if the specified index has depth set to 0 or the depth is greater than
139 /// the depth of this Merkle tree.
140 pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
141 if index.is_root() {
142 Err(MerkleError::DepthTooSmall(index.depth()))
143 } else if index.depth() > self.depth() {
144 Err(MerkleError::DepthTooBig(index.depth() as u64))
145 } else if index.depth() == self.depth() {
146 // the lookup in empty_hashes could fail only if empty_hashes were not built correctly
147 // by the constructor as we check the depth of the lookup above.
148 let leaf_pos = index.value();
149 let leaf = match self.get_leaf_node(leaf_pos) {
150 Some(word) => word.into(),
151 None => *EmptySubtreeRoots::entry(self.depth, index.depth()),
152 };
153 Ok(leaf)
154 } else {
155 Ok(self.get_branch_node(&index).parent())
156 }
157 }
158
159 /// Returns a value of the leaf at the specified index.
160 ///
161 /// # Errors
162 /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
163 pub fn get_leaf(&self, index: u64) -> Result<Word, MerkleError> {
164 let index = NodeIndex::new(self.depth, index)?;
165 Ok(self.get_node(index)?.into())
166 }
167
168 /// Returns a Merkle path from the node at the specified index to the root.
169 ///
170 /// The node itself is not included in the path.
171 ///
172 /// # Errors
173 /// Returns an error if the specified index has depth set to 0 or the depth is greater than
174 /// the depth of this Merkle tree.
175 pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
176 if index.is_root() {
177 return Err(MerkleError::DepthTooSmall(index.depth()));
178 } else if index.depth() > self.depth() {
179 return Err(MerkleError::DepthTooBig(index.depth() as u64));
180 }
181
182 let mut path = Vec::with_capacity(index.depth() as usize);
183 for _ in 0..index.depth() {
184 let is_right = index.is_value_odd();
185 index.move_up();
186 let BranchNode { left, right } = self.get_branch_node(&index);
187 let value = if is_right { left } else { right };
188 path.push(value);
189 }
190 Ok(MerklePath::new(path))
191 }
192
193 /// Return a Merkle path from the leaf at the specified index to the root.
194 ///
195 /// The leaf itself is not included in the path.
196 ///
197 /// # Errors
198 /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
199 pub fn get_leaf_path(&self, index: u64) -> Result<MerklePath, MerkleError> {
200 let index = NodeIndex::new(self.depth(), index)?;
201 self.get_path(index)
202 }
203
204 // ITERATORS
205 // --------------------------------------------------------------------------------------------
206
207 /// Returns an iterator over the leaves of this [SimpleSmt].
208 pub fn leaves(&self) -> impl Iterator<Item = (u64, &Word)> {
209 self.leaves.iter().map(|(i, w)| (*i, w))
210 }
211
212 /// Returns an iterator over the inner nodes of this Merkle tree.
213 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
214 self.branches.values().map(|e| InnerNodeInfo {
215 value: e.parent(),
216 left: e.left,
217 right: e.right,
218 })
219 }
220
221 // STATE MUTATORS
222 // --------------------------------------------------------------------------------------------
223
224 /// Updates value of the leaf at the specified index returning the old leaf value.
225 ///
226 /// This also recomputes all hashes between the leaf and the root, updating the root itself.
227 ///
228 /// # Errors
229 /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
230 pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<Word, MerkleError> {
231 // validate the index before modifying the structure
232 let idx = NodeIndex::new(self.depth(), index)?;
233
234 let old_value = self.insert_leaf_node(index, value).unwrap_or(Self::EMPTY_VALUE);
235
236 // if the old value and new value are the same, there is nothing to update
237 if value == old_value {
238 return Ok(value);
239 }
240
241 self.recompute_nodes_from_index_to_root(idx, RpoDigest::from(value));
242
243 Ok(old_value)
244 }
245
246 /// Inserts a subtree at the specified index. The depth at which the subtree is inserted is
247 /// computed as `self.depth() - subtree.depth()`.
248 ///
249 /// Returns the new root.
250 pub fn set_subtree(
251 &mut self,
252 subtree_insertion_index: u64,
253 subtree: SimpleSmt,
254 ) -> Result<RpoDigest, MerkleError> {
255 if subtree.depth() > self.depth() {
256 return Err(MerkleError::InvalidSubtreeDepth {
257 subtree_depth: subtree.depth(),
258 tree_depth: self.depth(),
259 });
260 }
261
262 // Verify that `subtree_insertion_index` is valid.
263 let subtree_root_insertion_depth = self.depth() - subtree.depth();
264 let subtree_root_index =
265 NodeIndex::new(subtree_root_insertion_depth, subtree_insertion_index)?;
266
267 // add leaves
268 // --------------
269
270 // The subtree's leaf indices live in their own context - i.e. a subtree of depth `d`. If we
271 // insert the subtree at `subtree_insertion_index = 0`, then the subtree leaf indices are
272 // valid as they are. However, consider what happens when we insert at
273 // `subtree_insertion_index = 1`. The first leaf of our subtree now will have index `2^d`;
274 // you can see it as there's a full subtree sitting on its left. In general, for
275 // `subtree_insertion_index = i`, there are `i` subtrees sitting before the subtree we want
276 // to insert, so we need to adjust all its leaves by `i * 2^d`.
277 let leaf_index_shift: u64 = subtree_insertion_index * 2_u64.pow(subtree.depth().into());
278 for (subtree_leaf_idx, leaf_value) in subtree.leaves() {
279 let new_leaf_idx = leaf_index_shift + subtree_leaf_idx;
280 debug_assert!(new_leaf_idx < 2_u64.pow(self.depth().into()));
281
282 self.insert_leaf_node(new_leaf_idx, *leaf_value);
283 }
284
285 // add subtree's branch nodes (which includes the root)
286 // --------------
287 for (branch_idx, branch_node) in subtree.branches {
288 let new_branch_idx = {
289 let new_depth = subtree_root_insertion_depth + branch_idx.depth();
290 let new_value = subtree_insertion_index * 2_u64.pow(branch_idx.depth().into())
291 + branch_idx.value();
292
293 NodeIndex::new(new_depth, new_value).expect("index guaranteed to be valid")
294 };
295
296 self.branches.insert(new_branch_idx, branch_node);
297 }
298
299 // recompute nodes starting from subtree root
300 // --------------
301 self.recompute_nodes_from_index_to_root(subtree_root_index, subtree.root);
302
303 Ok(self.root)
304 }
305
306 // HELPER METHODS
307 // --------------------------------------------------------------------------------------------
308
309 /// Recomputes the branch nodes (including the root) from `index` all the way to the root.
310 /// `node_hash_at_index` is the hash of the node stored at index.
311 fn recompute_nodes_from_index_to_root(
312 &mut self,
313 mut index: NodeIndex,
314 node_hash_at_index: RpoDigest,
315 ) {
316 let mut value = node_hash_at_index;
317 for _ in 0..index.depth() {
318 let is_right = index.is_value_odd();
319 index.move_up();
320 let BranchNode { left, right } = self.get_branch_node(&index);
321 let (left, right) = if is_right { (left, value) } else { (value, right) };
322 self.insert_branch_node(index, left, right);
323 value = Rpo256::merge(&[left, right]);
324 }
325 self.root = value;
326 }
327
328 fn get_leaf_node(&self, key: u64) -> Option<Word> {
329 self.leaves.get(&key).copied()
330 }
331
332 fn insert_leaf_node(&mut self, key: u64, node: Word) -> Option<Word> {
333 self.leaves.insert(key, node)
334 }
335
336 fn get_branch_node(&self, index: &NodeIndex) -> BranchNode {
337 self.branches.get(index).cloned().unwrap_or_else(|| {
338 let node = EmptySubtreeRoots::entry(self.depth, index.depth() + 1);
339 BranchNode { left: *node, right: *node }
340 })
341 }
342
343 fn insert_branch_node(&mut self, index: NodeIndex, left: RpoDigest, right: RpoDigest) {
344 let branch = BranchNode { left, right };
345 self.branches.insert(index, branch);
346 }
347}
348
349// BRANCH NODE
350// ================================================================================================
351
352#[derive(Debug, Default, Clone, PartialEq, Eq)]
353#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
354struct BranchNode {
355 left: RpoDigest,
356 right: RpoDigest,
357}
358
359impl BranchNode {
360 fn parent(&self) -> RpoDigest {
361 Rpo256::merge(&[self.left, self.right])
362 }
363}
364
365// TRY APPLY DIFF
366// ================================================================================================
367impl TryApplyDiff<RpoDigest, StoreNode> for SimpleSmt {
368 type Error = MerkleError;
369 type DiffType = MerkleTreeDelta;
370
371 fn try_apply(&mut self, diff: MerkleTreeDelta) -> Result<(), MerkleError> {
372 if diff.depth() != self.depth() {
373 return Err(MerkleError::InvalidDepth {
374 expected: self.depth(),
375 provided: diff.depth(),
376 });
377 }
378
379 for slot in diff.cleared_slots() {
380 self.update_leaf(*slot, Self::EMPTY_VALUE)?;
381 }
382
383 for (slot, value) in diff.updated_slots() {
384 self.update_leaf(*slot, *value)?;
385 }
386
387 Ok(())
388 }
389}