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