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