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