miden-crypto/src/merkle/simple_smt/mod.rs
src/merkle/simple_smt/mod.rsRaw
1use super::{
2 BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerkleTreeDelta,
3 NodeIndex, Rpo256, RpoDigest, StoreNode, TryApplyDiff, Vec, Word,
4};
5
6#[cfg(test)]
7mod tests;
8
9// SPARSE MERKLE TREE
10// ================================================================================================
11
12/// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction.
13///
14/// The root of the tree is recomputed on each new leaf update.
15#[derive(Debug, Clone, PartialEq, Eq)]
16#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
17pub struct SimpleSmt {
18 depth: u8,
19 root: RpoDigest,
20 leaves: BTreeMap<u64, Word>,
21 branches: BTreeMap<NodeIndex, BranchNode>,
22}
23
24impl SimpleSmt {
25 // CONSTANTS
26 // --------------------------------------------------------------------------------------------
27
28 /// Minimum supported depth.
29 pub const MIN_DEPTH: u8 = 1;
30
31 /// Maximum supported depth.
32 pub const MAX_DEPTH: u8 = 64;
33
34 /// Value of an empty leaf.
35 pub const EMPTY_VALUE: Word = super::EMPTY_WORD;
36
37 // CONSTRUCTORS
38 // --------------------------------------------------------------------------------------------
39
40 /// Returns a new [SimpleSmt] instantiated with the specified depth.
41 ///
42 /// All leaves in the returned tree are set to [ZERO; 4].
43 ///
44 /// # Errors
45 /// Returns an error if the depth is 0 or is greater than 64.
46 pub fn new(depth: u8) -> Result<Self, MerkleError> {
47 // validate the range of the depth.
48 if depth < Self::MIN_DEPTH {
49 return Err(MerkleError::DepthTooSmall(depth));
50 } else if Self::MAX_DEPTH < depth {
51 return Err(MerkleError::DepthTooBig(depth as u64));
52 }
53
54 let root = *EmptySubtreeRoots::entry(depth, 0);
55
56 Ok(Self {
57 root,
58 depth,
59 leaves: BTreeMap::new(),
60 branches: BTreeMap::new(),
61 })
62 }
63
64 /// Returns a new [SimpleSmt] instantiated with the specified depth and with leaves
65 /// set as specified by the provided entries.
66 ///
67 /// All leaves omitted from the entries list are set to [ZERO; 4].
68 ///
69 /// # Errors
70 /// Returns an error if:
71 /// - If the depth is 0 or is greater than 64.
72 /// - The number of entries exceeds the maximum tree capacity, that is 2^{depth}.
73 /// - The provided entries contain multiple values for the same key.
74 pub fn with_leaves<R, I>(depth: u8, entries: R) -> Result<Self, MerkleError>
75 where
76 R: IntoIterator<IntoIter = I>,
77 I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
78 {
79 // create an empty tree
80 let mut tree = Self::new(depth)?;
81
82 // check if the number of leaves can be accommodated by the tree's depth; we use a min
83 // depth of 63 because we consider passing in a vector of size 2^64 infeasible.
84 let entries = entries.into_iter();
85 let max = 1 << tree.depth.min(63);
86 if entries.len() > max {
87 return Err(MerkleError::InvalidNumEntries(max, entries.len()));
88 }
89
90 // append leaves to the tree returning an error if a duplicate entry for the same key
91 // is found
92 let mut empty_entries = BTreeSet::new();
93 for (key, value) in entries {
94 let old_value = tree.update_leaf(key, value)?;
95 if old_value != Self::EMPTY_VALUE || empty_entries.contains(&key) {
96 return Err(MerkleError::DuplicateValuesForIndex(key));
97 }
98 // if we've processed an empty entry, add the key to the set of empty entry keys, and
99 // if this key was already in the set, return an error
100 if value == Self::EMPTY_VALUE && !empty_entries.insert(key) {
101 return Err(MerkleError::DuplicateValuesForIndex(key));
102 }
103 }
104 Ok(tree)
105 }
106
107 /// Wrapper around [`SimpleSmt::with_leaves`] which inserts leaves at contiguous indices
108 /// starting at index 0.
109 pub fn with_contiguous_leaves<R, I>(depth: u8, entries: R) -> Result<Self, MerkleError>
110 where
111 R: IntoIterator<IntoIter = I>,
112 I: Iterator<Item = Word> + ExactSizeIterator,
113 {
114 Self::with_leaves(
115 depth,
116 entries
117 .into_iter()
118 .enumerate()
119 .map(|(idx, word)| (idx.try_into().expect("tree max depth is 2^8"), word)),
120 )
121 }
122
123 // PUBLIC ACCESSORS
124 // --------------------------------------------------------------------------------------------
125
126 /// Returns the root of this Merkle tree.
127 pub const fn root(&self) -> RpoDigest {
128 self.root
129 }
130
131 /// Returns the depth of this Merkle tree.
132 pub const fn depth(&self) -> u8 {
133 self.depth
134 }
135
136 /// Returns a node at the specified index.
137 ///
138 /// # Errors
139 /// Returns an error if the specified index has depth set to 0 or the depth is greater than
140 /// the depth of this Merkle tree.
141 pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
142 if index.is_root() {
143 Err(MerkleError::DepthTooSmall(index.depth()))
144 } else if index.depth() > self.depth() {
145 Err(MerkleError::DepthTooBig(index.depth() as u64))
146 } else if index.depth() == self.depth() {
147 // the lookup in empty_hashes could fail only if empty_hashes were not built correctly
148 // by the constructor as we check the depth of the lookup above.
149 let leaf_pos = index.value();
150 let leaf = match self.get_leaf_node(leaf_pos) {
151 Some(word) => word.into(),
152 None => *EmptySubtreeRoots::entry(self.depth, index.depth()),
153 };
154 Ok(leaf)
155 } else {
156 Ok(self.get_branch_node(&index).parent())
157 }
158 }
159
160 /// Returns a value of the leaf at the specified index.
161 ///
162 /// # Errors
163 /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
164 pub fn get_leaf(&self, index: u64) -> Result<Word, MerkleError> {
165 let index = NodeIndex::new(self.depth, index)?;
166 Ok(self.get_node(index)?.into())
167 }
168
169 /// Returns a Merkle path from the node at the specified index to the root.
170 ///
171 /// The node itself is not included in the path.
172 ///
173 /// # Errors
174 /// Returns an error if the specified index has depth set to 0 or the depth is greater than
175 /// the depth of this Merkle tree.
176 pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
177 if index.is_root() {
178 return Err(MerkleError::DepthTooSmall(index.depth()));
179 } else if index.depth() > self.depth() {
180 return Err(MerkleError::DepthTooBig(index.depth() as u64));
181 }
182
183 let mut path = Vec::with_capacity(index.depth() as usize);
184 for _ in 0..index.depth() {
185 let is_right = index.is_value_odd();
186 index.move_up();
187 let BranchNode { left, right } = self.get_branch_node(&index);
188 let value = if is_right { left } else { right };
189 path.push(value);
190 }
191 Ok(MerklePath::new(path))
192 }
193
194 /// Return a Merkle path from the leaf at the specified index to the root.
195 ///
196 /// The leaf itself is not included in the path.
197 ///
198 /// # Errors
199 /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
200 pub fn get_leaf_path(&self, index: u64) -> Result<MerklePath, MerkleError> {
201 let index = NodeIndex::new(self.depth(), index)?;
202 self.get_path(index)
203 }
204
205 // ITERATORS
206 // --------------------------------------------------------------------------------------------
207
208 /// Returns an iterator over the leaves of this [SimpleSmt].
209 pub fn leaves(&self) -> impl Iterator<Item = (u64, &Word)> {
210 self.leaves.iter().map(|(i, w)| (*i, w))
211 }
212
213 /// Returns an iterator over the inner nodes of this Merkle tree.
214 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
215 self.branches.values().map(|e| InnerNodeInfo {
216 value: e.parent(),
217 left: e.left,
218 right: e.right,
219 })
220 }
221
222 // STATE MUTATORS
223 // --------------------------------------------------------------------------------------------
224
225 /// Updates value of the leaf at the specified index returning the old leaf value.
226 ///
227 /// This also recomputes all hashes between the leaf and the root, updating the root itself.
228 ///
229 /// # Errors
230 /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
231 pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<Word, MerkleError> {
232 // validate the index before modifying the structure
233 let mut idx = NodeIndex::new(self.depth(), index)?;
234
235 let old_value = self.insert_leaf_node(index, value).unwrap_or(Self::EMPTY_VALUE);
236
237 // if the old value and new value are the same, there is nothing to update
238 if value == old_value {
239 return Ok(value);
240 }
241
242 let mut value = RpoDigest::from(value);
243 for _ in 0..idx.depth() {
244 let is_right = idx.is_value_odd();
245 idx.move_up();
246 let BranchNode { left, right } = self.get_branch_node(&idx);
247 let (left, right) = if is_right { (left, value) } else { (value, right) };
248 self.insert_branch_node(idx, left, right);
249 value = Rpo256::merge(&[left, right]);
250 }
251 self.root = value;
252 Ok(old_value)
253 }
254
255 // HELPER METHODS
256 // --------------------------------------------------------------------------------------------
257
258 fn get_leaf_node(&self, key: u64) -> Option<Word> {
259 self.leaves.get(&key).copied()
260 }
261
262 fn insert_leaf_node(&mut self, key: u64, node: Word) -> Option<Word> {
263 self.leaves.insert(key, node)
264 }
265
266 fn get_branch_node(&self, index: &NodeIndex) -> BranchNode {
267 self.branches.get(index).cloned().unwrap_or_else(|| {
268 let node = EmptySubtreeRoots::entry(self.depth, index.depth() + 1);
269 BranchNode { left: *node, right: *node }
270 })
271 }
272
273 fn insert_branch_node(&mut self, index: NodeIndex, left: RpoDigest, right: RpoDigest) {
274 let branch = BranchNode { left, right };
275 self.branches.insert(index, branch);
276 }
277}
278
279// BRANCH NODE
280// ================================================================================================
281
282#[derive(Debug, Default, Clone, PartialEq, Eq)]
283#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
284struct BranchNode {
285 left: RpoDigest,
286 right: RpoDigest,
287}
288
289impl BranchNode {
290 fn parent(&self) -> RpoDigest {
291 Rpo256::merge(&[self.left, self.right])
292 }
293}
294
295// TRY APPLY DIFF
296// ================================================================================================
297impl TryApplyDiff<RpoDigest, StoreNode> for SimpleSmt {
298 type Error = MerkleError;
299 type DiffType = MerkleTreeDelta;
300
301 fn try_apply(&mut self, diff: MerkleTreeDelta) -> Result<(), MerkleError> {
302 if diff.depth() != self.depth() {
303 return Err(MerkleError::InvalidDepth {
304 expected: self.depth(),
305 provided: diff.depth(),
306 });
307 }
308
309 for slot in diff.cleared_slots() {
310 self.update_leaf(*slot, Self::EMPTY_VALUE)?;
311 }
312
313 for (slot, value) in diff.updated_slots() {
314 self.update_leaf(*slot, *value)?;
315 }
316
317 Ok(())
318 }
319}