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