miden-crypto/src/merkle/smt/full/mod.rs
src/merkle/smt/full/mod.rsRaw
1use winter_math::StarkField;
2
3use crate::hash::rpo::Rpo256;
4use crate::merkle::{EmptySubtreeRoots, InnerNodeInfo};
5use crate::utils::collections::{BTreeMap, BTreeSet};
6use crate::{Felt, EMPTY_WORD};
7
8use super::{
9 InnerNode, LeafIndex, MerkleError, MerklePath, NodeIndex, RpoDigest, SparseMerkleTree, Word,
10};
11
12mod error;
13pub use error::{SmtLeafError, SmtProofError};
14
15mod leaf;
16pub use leaf::SmtLeaf;
17
18mod proof;
19pub use proof::SmtProof;
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}