miden-crypto/src/merkle/smt/full/mod.rs
src/merkle/smt/full/mod.rsRaw
1use super::{
2 EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, MerklePath,
3 NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD,
4};
5use alloc::collections::{BTreeMap, BTreeSet};
6
7mod error;
8pub use error::{SmtLeafError, SmtProofError};
9
10mod leaf;
11pub use leaf::SmtLeaf;
12
13mod proof;
14pub use proof::SmtProof;
15
16#[cfg(test)]
17mod tests;
18
19// CONSTANTS
20// ================================================================================================
21
22pub const SMT_DEPTH: u8 = 64;
23
24// SMT
25// ================================================================================================
26
27/// Sparse Merkle tree mapping 256-bit keys to 256-bit values. Both keys and values are represented
28/// by 4 field elements.
29///
30/// All leaves sit at depth 64. The most significant element of the key is used to identify the leaf to
31/// which the key maps.
32///
33/// A leaf is either empty, or holds one or more key-value pairs. An empty leaf hashes to the empty
34/// word. Otherwise, a leaf hashes to the hash of its key-value pairs, ordered by key first, value
35/// second.
36#[derive(Debug, Clone, PartialEq, Eq)]
37#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
38pub struct Smt {
39 root: RpoDigest,
40 leaves: BTreeMap<u64, SmtLeaf>,
41 inner_nodes: BTreeMap<NodeIndex, InnerNode>,
42}
43
44impl Smt {
45 // CONSTANTS
46 // --------------------------------------------------------------------------------------------
47 /// The default value used to compute the hash of empty leaves
48 pub const EMPTY_VALUE: Word = <Self as SparseMerkleTree<SMT_DEPTH>>::EMPTY_VALUE;
49
50 // CONSTRUCTORS
51 // --------------------------------------------------------------------------------------------
52
53 /// Returns a new [Smt].
54 ///
55 /// All leaves in the returned tree are set to [Self::EMPTY_VALUE].
56 pub fn new() -> Self {
57 let root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
58
59 Self {
60 root,
61 leaves: BTreeMap::new(),
62 inner_nodes: BTreeMap::new(),
63 }
64 }
65
66 /// Returns a new [Smt] instantiated with leaves set as specified by the provided entries.
67 ///
68 /// All leaves omitted from the entries list are set to [Self::EMPTY_VALUE].
69 ///
70 /// # Errors
71 /// Returns an error if the provided entries contain multiple values for the same key.
72 pub fn with_entries(
73 entries: impl IntoIterator<Item = (RpoDigest, Word)>,
74 ) -> Result<Self, MerkleError> {
75 // create an empty tree
76 let mut tree = Self::new();
77
78 // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
79 // entries with the empty value need additional tracking.
80 let mut key_set_to_zero = BTreeSet::new();
81
82 for (key, value) in entries {
83 let old_value = tree.insert(key, value);
84
85 if old_value != EMPTY_WORD || key_set_to_zero.contains(&key) {
86 return Err(MerkleError::DuplicateValuesForIndex(
87 LeafIndex::<SMT_DEPTH>::from(key).value(),
88 ));
89 }
90
91 if value == EMPTY_WORD {
92 key_set_to_zero.insert(key);
93 };
94 }
95 Ok(tree)
96 }
97
98 // PUBLIC ACCESSORS
99 // --------------------------------------------------------------------------------------------
100
101 /// Returns the depth of the tree
102 pub const fn depth(&self) -> u8 {
103 SMT_DEPTH
104 }
105
106 /// Returns the root of the tree
107 pub fn root(&self) -> RpoDigest {
108 <Self as SparseMerkleTree<SMT_DEPTH>>::root(self)
109 }
110
111 /// Returns the leaf to which `key` maps
112 pub fn get_leaf(&self, key: &RpoDigest) -> SmtLeaf {
113 <Self as SparseMerkleTree<SMT_DEPTH>>::get_leaf(self, key)
114 }
115
116 /// Returns the value associated with `key`
117 pub fn get_value(&self, key: &RpoDigest) -> Word {
118 let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
119
120 match self.leaves.get(&leaf_pos) {
121 Some(leaf) => leaf.get_value(key).unwrap_or_default(),
122 None => EMPTY_WORD,
123 }
124 }
125
126 /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
127 /// path to the leaf, as well as the leaf itself.
128 pub fn open(&self, key: &RpoDigest) -> SmtProof {
129 <Self as SparseMerkleTree<SMT_DEPTH>>::open(self, key)
130 }
131
132 // ITERATORS
133 // --------------------------------------------------------------------------------------------
134
135 /// Returns an iterator over the leaves of this [Smt].
136 pub fn leaves(&self) -> impl Iterator<Item = (LeafIndex<SMT_DEPTH>, &SmtLeaf)> {
137 self.leaves
138 .iter()
139 .map(|(leaf_index, leaf)| (LeafIndex::new_max_depth(*leaf_index), leaf))
140 }
141
142 /// Returns an iterator over the key-value pairs of this [Smt].
143 pub fn entries(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
144 self.leaves().flat_map(|(_, leaf)| leaf.entries())
145 }
146
147 /// Returns an iterator over the inner nodes of this [Smt].
148 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
149 self.inner_nodes.values().map(|e| InnerNodeInfo {
150 value: e.hash(),
151 left: e.left,
152 right: e.right,
153 })
154 }
155
156 // STATE MUTATORS
157 // --------------------------------------------------------------------------------------------
158
159 /// Inserts a value at the specified key, returning the previous value associated with that key.
160 /// Recall that by definition, any key that hasn't been updated is associated with
161 /// [`Self::EMPTY_VALUE`].
162 ///
163 /// This also recomputes all hashes between the leaf (associated with the key) and the root,
164 /// updating the root itself.
165 pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word {
166 <Self as SparseMerkleTree<SMT_DEPTH>>::insert(self, key, value)
167 }
168
169 // HELPERS
170 // --------------------------------------------------------------------------------------------
171
172 /// Inserts `value` at leaf index pointed to by `key`. `value` is guaranteed to not be the empty
173 /// value, such that this is indeed an insertion.
174 fn perform_insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
175 debug_assert_ne!(value, Self::EMPTY_VALUE);
176
177 let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);
178
179 match self.leaves.get_mut(&leaf_index.value()) {
180 Some(leaf) => leaf.insert(key, value),
181 None => {
182 self.leaves.insert(leaf_index.value(), SmtLeaf::Single((key, value)));
183
184 None
185 }
186 }
187 }
188
189 /// Removes key-value pair at leaf index pointed to by `key` if it exists.
190 fn perform_remove(&mut self, key: RpoDigest) -> Option<Word> {
191 let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);
192
193 if let Some(leaf) = self.leaves.get_mut(&leaf_index.value()) {
194 let (old_value, is_empty) = leaf.remove(key);
195 if is_empty {
196 self.leaves.remove(&leaf_index.value());
197 }
198 old_value
199 } else {
200 // there's nothing stored at the leaf; nothing to update
201 None
202 }
203 }
204}
205
206impl SparseMerkleTree<SMT_DEPTH> for Smt {
207 type Key = RpoDigest;
208 type Value = Word;
209 type Leaf = SmtLeaf;
210 type Opening = SmtProof;
211
212 const EMPTY_VALUE: Self::Value = EMPTY_WORD;
213
214 fn root(&self) -> RpoDigest {
215 self.root
216 }
217
218 fn set_root(&mut self, root: RpoDigest) {
219 self.root = root;
220 }
221
222 fn get_inner_node(&self, index: NodeIndex) -> InnerNode {
223 self.inner_nodes.get(&index).cloned().unwrap_or_else(|| {
224 let node = EmptySubtreeRoots::entry(SMT_DEPTH, index.depth() + 1);
225
226 InnerNode { left: *node, right: *node }
227 })
228 }
229
230 fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) {
231 self.inner_nodes.insert(index, inner_node);
232 }
233
234 fn remove_inner_node(&mut self, index: NodeIndex) {
235 let _ = self.inner_nodes.remove(&index);
236 }
237
238 fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value> {
239 // inserting an `EMPTY_VALUE` is equivalent to removing any value associated with `key`
240 if value != Self::EMPTY_VALUE {
241 self.perform_insert(key, value)
242 } else {
243 self.perform_remove(key)
244 }
245 }
246
247 fn get_leaf(&self, key: &RpoDigest) -> Self::Leaf {
248 let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
249
250 match self.leaves.get(&leaf_pos) {
251 Some(leaf) => leaf.clone(),
252 None => SmtLeaf::new_empty(key.into()),
253 }
254 }
255
256 fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest {
257 leaf.hash()
258 }
259
260 fn key_to_leaf_index(key: &RpoDigest) -> LeafIndex<SMT_DEPTH> {
261 let most_significant_felt = key[3];
262 LeafIndex::new_max_depth(most_significant_felt.as_int())
263 }
264
265 fn path_and_leaf_to_opening(path: MerklePath, leaf: SmtLeaf) -> SmtProof {
266 SmtProof::new_unchecked(path, leaf)
267 }
268}
269
270impl Default for Smt {
271 fn default() -> Self {
272 Self::new()
273 }
274}
275
276// CONVERSIONS
277// ================================================================================================
278
279impl From<Word> for LeafIndex<SMT_DEPTH> {
280 fn from(value: Word) -> Self {
281 // We use the most significant `Felt` of a `Word` as the leaf index.
282 Self::new_max_depth(value[3].as_int())
283 }
284}
285
286impl From<RpoDigest> for LeafIndex<SMT_DEPTH> {
287 fn from(value: RpoDigest) -> Self {
288 Word::from(value).into()
289 }
290}
291
292impl From<&RpoDigest> for LeafIndex<SMT_DEPTH> {
293 fn from(value: &RpoDigest) -> Self {
294 Word::from(value).into()
295 }
296}