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