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