miden-crypto/src/merkle/smt/simple/mod.rs
src/merkle/smt/simple/mod.rsRaw
1use crate::{
2 merkle::{EmptySubtreeRoots, InnerNodeInfo, ValuePath},
3 EMPTY_WORD,
4};
5
6use super::{
7 InnerNode, LeafIndex, MerkleError, NodeIndex, RpoDigest, SparseMerkleTree, Word, SMT_MAX_DEPTH,
8 SMT_MIN_DEPTH,
9};
10use crate::utils::collections::{BTreeMap, BTreeSet};
11
12#[cfg(test)]
13mod tests;
14
15// SPARSE MERKLE TREE
16// ================================================================================================
17
18/// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction.
19///
20/// The root of the tree is recomputed on each new leaf update.
21#[derive(Debug, Clone, PartialEq, Eq)]
22#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
23pub struct SimpleSmt<const DEPTH: u8> {
24 root: RpoDigest,
25 leaves: BTreeMap<u64, Word>,
26 inner_nodes: BTreeMap<NodeIndex, InnerNode>,
27}
28
29impl<const DEPTH: u8> SimpleSmt<DEPTH> {
30 // CONSTANTS
31 // --------------------------------------------------------------------------------------------
32
33 /// The default value used to compute the hash of empty leaves
34 pub const EMPTY_VALUE: Word = <Self as SparseMerkleTree<DEPTH>>::EMPTY_VALUE;
35
36 // CONSTRUCTORS
37 // --------------------------------------------------------------------------------------------
38
39 /// Returns a new [SimpleSmt].
40 ///
41 /// All leaves in the returned tree are set to [ZERO; 4].
42 ///
43 /// # Errors
44 /// Returns an error if DEPTH is 0 or is greater than 64.
45 pub fn new() -> Result<Self, MerkleError> {
46 // validate the range of the depth.
47 if DEPTH < SMT_MIN_DEPTH {
48 return Err(MerkleError::DepthTooSmall(DEPTH));
49 } else if SMT_MAX_DEPTH < DEPTH {
50 return Err(MerkleError::DepthTooBig(DEPTH as u64));
51 }
52
53 let root = *EmptySubtreeRoots::entry(DEPTH, 0);
54
55 Ok(Self {
56 root,
57 leaves: BTreeMap::new(),
58 inner_nodes: BTreeMap::new(),
59 })
60 }
61
62 /// Returns a new [SimpleSmt] instantiated with leaves set as specified by the provided entries.
63 ///
64 /// All leaves omitted from the entries list are set to [ZERO; 4].
65 ///
66 /// # Errors
67 /// Returns an error if:
68 /// - If the depth is 0 or is greater than 64.
69 /// - The number of entries exceeds the maximum tree capacity, that is 2^{depth}.
70 /// - The provided entries contain multiple values for the same key.
71 pub fn with_leaves(
72 entries: impl IntoIterator<Item = (u64, Word)>,
73 ) -> Result<Self, MerkleError> {
74 // create an empty tree
75 let mut tree = Self::new()?;
76
77 // compute the max number of entries. We use an upper bound of depth 63 because we consider
78 // passing in a vector of size 2^64 infeasible.
79 let max_num_entries = 2_usize.pow(DEPTH.min(63).into());
80
81 // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
82 // entries with the empty value need additional tracking.
83 let mut key_set_to_zero = BTreeSet::new();
84
85 for (idx, (key, value)) in entries.into_iter().enumerate() {
86 if idx >= max_num_entries {
87 return Err(MerkleError::InvalidNumEntries(max_num_entries));
88 }
89
90 let old_value = tree.insert(LeafIndex::<DEPTH>::new(key)?, value);
91
92 if old_value != Self::EMPTY_VALUE || key_set_to_zero.contains(&key) {
93 return Err(MerkleError::DuplicateValuesForIndex(key));
94 }
95
96 if value == Self::EMPTY_VALUE {
97 key_set_to_zero.insert(key);
98 };
99 }
100 Ok(tree)
101 }
102
103 /// Wrapper around [`SimpleSmt::with_leaves`] which inserts leaves at contiguous indices
104 /// starting at index 0.
105 pub fn with_contiguous_leaves(
106 entries: impl IntoIterator<Item = Word>,
107 ) -> Result<Self, MerkleError> {
108 Self::with_leaves(
109 entries
110 .into_iter()
111 .enumerate()
112 .map(|(idx, word)| (idx.try_into().expect("tree max depth is 2^8"), word)),
113 )
114 }
115
116 // PUBLIC ACCESSORS
117 // --------------------------------------------------------------------------------------------
118
119 /// Returns the depth of the tree
120 pub const fn depth(&self) -> u8 {
121 DEPTH
122 }
123
124 /// Returns the root of the tree
125 pub fn root(&self) -> RpoDigest {
126 <Self as SparseMerkleTree<DEPTH>>::root(self)
127 }
128
129 /// Returns the leaf at the specified index.
130 pub fn get_leaf(&self, key: &LeafIndex<DEPTH>) -> Word {
131 <Self as SparseMerkleTree<DEPTH>>::get_leaf(self, key)
132 }
133
134 /// Returns a node at the specified index.
135 ///
136 /// # Errors
137 /// Returns an error if the specified index has depth set to 0 or the depth is greater than
138 /// the depth of this Merkle tree.
139 pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
140 if index.is_root() {
141 Err(MerkleError::DepthTooSmall(index.depth()))
142 } else if index.depth() > DEPTH {
143 Err(MerkleError::DepthTooBig(index.depth() as u64))
144 } else if index.depth() == DEPTH {
145 let leaf = self.get_leaf(&LeafIndex::<DEPTH>::try_from(index)?);
146
147 Ok(leaf.into())
148 } else {
149 Ok(self.get_inner_node(index).hash())
150 }
151 }
152
153 /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
154 /// path to the leaf, as well as the leaf itself.
155 pub fn open(&self, key: &LeafIndex<DEPTH>) -> ValuePath {
156 <Self as SparseMerkleTree<DEPTH>>::open(self, key)
157 }
158
159 // ITERATORS
160 // --------------------------------------------------------------------------------------------
161
162 /// Returns an iterator over the leaves of this [SimpleSmt].
163 pub fn leaves(&self) -> impl Iterator<Item = (u64, &Word)> {
164 self.leaves.iter().map(|(i, w)| (*i, w))
165 }
166
167 /// Returns an iterator over the inner nodes of this [SimpleSmt].
168 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
169 self.inner_nodes.values().map(|e| InnerNodeInfo {
170 value: e.hash(),
171 left: e.left,
172 right: e.right,
173 })
174 }
175
176 // STATE MUTATORS
177 // --------------------------------------------------------------------------------------------
178
179 /// Inserts a value at the specified key, returning the previous value associated with that key.
180 /// Recall that by definition, any key that hasn't been updated is associated with
181 /// [`EMPTY_WORD`].
182 ///
183 /// This also recomputes all hashes between the leaf (associated with the key) and the root,
184 /// updating the root itself.
185 pub fn insert(&mut self, key: LeafIndex<DEPTH>, value: Word) -> Word {
186 <Self as SparseMerkleTree<DEPTH>>::insert(self, key, value)
187 }
188
189 /// Inserts a subtree at the specified index. The depth at which the subtree is inserted is
190 /// computed as `DEPTH - SUBTREE_DEPTH`.
191 ///
192 /// Returns the new root.
193 pub fn set_subtree<const SUBTREE_DEPTH: u8>(
194 &mut self,
195 subtree_insertion_index: u64,
196 subtree: SimpleSmt<SUBTREE_DEPTH>,
197 ) -> Result<RpoDigest, MerkleError> {
198 if SUBTREE_DEPTH > DEPTH {
199 return Err(MerkleError::InvalidSubtreeDepth {
200 subtree_depth: SUBTREE_DEPTH,
201 tree_depth: DEPTH,
202 });
203 }
204
205 // Verify that `subtree_insertion_index` is valid.
206 let subtree_root_insertion_depth = DEPTH - SUBTREE_DEPTH;
207 let subtree_root_index =
208 NodeIndex::new(subtree_root_insertion_depth, subtree_insertion_index)?;
209
210 // add leaves
211 // --------------
212
213 // The subtree's leaf indices live in their own context - i.e. a subtree of depth `d`. If we
214 // insert the subtree at `subtree_insertion_index = 0`, then the subtree leaf indices are
215 // valid as they are. However, consider what happens when we insert at
216 // `subtree_insertion_index = 1`. The first leaf of our subtree now will have index `2^d`;
217 // you can see it as there's a full subtree sitting on its left. In general, for
218 // `subtree_insertion_index = i`, there are `i` subtrees sitting before the subtree we want
219 // to insert, so we need to adjust all its leaves by `i * 2^d`.
220 let leaf_index_shift: u64 = subtree_insertion_index * 2_u64.pow(SUBTREE_DEPTH.into());
221 for (subtree_leaf_idx, leaf_value) in subtree.leaves() {
222 let new_leaf_idx = leaf_index_shift + subtree_leaf_idx;
223 debug_assert!(new_leaf_idx < 2_u64.pow(DEPTH.into()));
224
225 self.leaves.insert(new_leaf_idx, *leaf_value);
226 }
227
228 // add subtree's branch nodes (which includes the root)
229 // --------------
230 for (branch_idx, branch_node) in subtree.inner_nodes {
231 let new_branch_idx = {
232 let new_depth = subtree_root_insertion_depth + branch_idx.depth();
233 let new_value = subtree_insertion_index * 2_u64.pow(branch_idx.depth().into())
234 + branch_idx.value();
235
236 NodeIndex::new(new_depth, new_value).expect("index guaranteed to be valid")
237 };
238
239 self.inner_nodes.insert(new_branch_idx, branch_node);
240 }
241
242 // recompute nodes starting from subtree root
243 // --------------
244 self.recompute_nodes_from_index_to_root(subtree_root_index, subtree.root);
245
246 Ok(self.root)
247 }
248}
249
250impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
251 type Key = LeafIndex<DEPTH>;
252 type Value = Word;
253 type Leaf = Word;
254 type Opening = ValuePath;
255
256 const EMPTY_VALUE: Self::Value = EMPTY_WORD;
257
258 fn root(&self) -> RpoDigest {
259 self.root
260 }
261
262 fn set_root(&mut self, root: RpoDigest) {
263 self.root = root;
264 }
265
266 fn get_inner_node(&self, index: NodeIndex) -> InnerNode {
267 self.inner_nodes.get(&index).cloned().unwrap_or_else(|| {
268 let node = EmptySubtreeRoots::entry(DEPTH, index.depth() + 1);
269
270 InnerNode { left: *node, right: *node }
271 })
272 }
273
274 fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) {
275 self.inner_nodes.insert(index, inner_node);
276 }
277
278 fn remove_inner_node(&mut self, index: NodeIndex) {
279 let _ = self.inner_nodes.remove(&index);
280 }
281
282 fn insert_value(&mut self, key: LeafIndex<DEPTH>, value: Word) -> Option<Word> {
283 self.leaves.insert(key.value(), value)
284 }
285
286 fn get_leaf(&self, key: &LeafIndex<DEPTH>) -> Word {
287 // the lookup in empty_hashes could fail only if empty_hashes were not built correctly
288 // by the constructor as we check the depth of the lookup above.
289 let leaf_pos = key.value();
290
291 match self.leaves.get(&leaf_pos) {
292 Some(word) => *word,
293 None => Word::from(*EmptySubtreeRoots::entry(DEPTH, DEPTH)),
294 }
295 }
296
297 fn hash_leaf(leaf: &Word) -> RpoDigest {
298 // `SimpleSmt` takes the leaf value itself as the hash
299 leaf.into()
300 }
301
302 fn key_to_leaf_index(key: &LeafIndex<DEPTH>) -> LeafIndex<DEPTH> {
303 *key
304 }
305}