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