miden-crypto/src/merkle/simple_smt/mod.rs
src/merkle/simple_smt/mod.rsRaw
1use super::{
2 BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex,
3 Rpo256, RpoDigest, Vec, Word, EMPTY_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: Word,
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].into();
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::DuplicateValuesForKey(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::DuplicateValuesForKey(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) -> Word {
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<Word, 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(self
133 .get_leaf_node(index.value())
134 .unwrap_or_else(|| self.empty_hashes[index.depth() as usize].into()))
135 } else {
136 Ok(self.get_branch_node(&index).parent().into())
137 }
138 }
139
140 /// Returns a value of the leaf at the specified index.
141 ///
142 /// # Errors
143 /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
144 pub fn get_leaf(&self, index: u64) -> Result<Word, MerkleError> {
145 let index = NodeIndex::new(self.depth, index)?;
146 self.get_node(index)
147 }
148
149 /// Returns a Merkle path from the node at the specified index to the root.
150 ///
151 /// The node itself is not included in the path.
152 ///
153 /// # Errors
154 /// Returns an error if the specified index has depth set to 0 or the depth is greater than
155 /// the depth of this Merkle tree.
156 pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
157 if index.is_root() {
158 return Err(MerkleError::DepthTooSmall(index.depth()));
159 } else if index.depth() > self.depth() {
160 return Err(MerkleError::DepthTooBig(index.depth() as u64));
161 }
162
163 let mut path = Vec::with_capacity(index.depth() as usize);
164 for _ in 0..index.depth() {
165 let is_right = index.is_value_odd();
166 index.move_up();
167 let BranchNode { left, right } = self.get_branch_node(&index);
168 let value = if is_right { left } else { right };
169 path.push(*value);
170 }
171 Ok(path.into())
172 }
173
174 /// Return a Merkle path from the leaf at the specified index to the root.
175 ///
176 /// The leaf itself is not included in the path.
177 ///
178 /// # Errors
179 /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
180 pub fn get_leaf_path(&self, index: u64) -> Result<MerklePath, MerkleError> {
181 let index = NodeIndex::new(self.depth(), index)?;
182 self.get_path(index)
183 }
184
185 // ITERATORS
186 // --------------------------------------------------------------------------------------------
187
188 /// Returns an iterator over the leaves of this [SimpleSmt].
189 pub fn leaves(&self) -> impl Iterator<Item = (u64, &Word)> {
190 self.leaves.iter().map(|(i, w)| (*i, w))
191 }
192
193 /// Returns an iterator over the inner nodes of this Merkle tree.
194 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
195 self.branches.values().map(|e| InnerNodeInfo {
196 value: e.parent().into(),
197 left: e.left.into(),
198 right: e.right.into(),
199 })
200 }
201
202 // STATE MUTATORS
203 // --------------------------------------------------------------------------------------------
204
205 /// Updates value of the leaf at the specified index returning the old leaf value.
206 ///
207 /// This also recomputes all hashes between the leaf and the root, updating the root itself.
208 ///
209 /// # Errors
210 /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
211 pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<Word, MerkleError> {
212 let old_value = self.insert_leaf_node(index, value).unwrap_or(EMPTY_WORD);
213
214 // if the old value and new value are the same, there is nothing to update
215 if value == old_value {
216 return Ok(value);
217 }
218
219 let mut index = NodeIndex::new(self.depth(), index)?;
220 let mut value = RpoDigest::from(value);
221 for _ in 0..index.depth() {
222 let is_right = index.is_value_odd();
223 index.move_up();
224 let BranchNode { left, right } = self.get_branch_node(&index);
225 let (left, right) = if is_right { (left, value) } else { (value, right) };
226 self.insert_branch_node(index, left, right);
227 value = Rpo256::merge(&[left, right]);
228 }
229 self.root = value.into();
230 Ok(old_value)
231 }
232
233 // HELPER METHODS
234 // --------------------------------------------------------------------------------------------
235
236 fn get_leaf_node(&self, key: u64) -> Option<Word> {
237 self.leaves.get(&key).copied()
238 }
239
240 fn insert_leaf_node(&mut self, key: u64, node: Word) -> Option<Word> {
241 self.leaves.insert(key, node)
242 }
243
244 fn get_branch_node(&self, index: &NodeIndex) -> BranchNode {
245 self.branches.get(index).cloned().unwrap_or_else(|| {
246 let node = self.empty_hashes[index.depth() as usize + 1];
247 BranchNode {
248 left: node,
249 right: node,
250 }
251 })
252 }
253
254 fn insert_branch_node(&mut self, index: NodeIndex, left: RpoDigest, right: RpoDigest) {
255 let branch = BranchNode { left, right };
256 self.branches.insert(index, branch);
257 }
258}
259
260// BRANCH NODE
261// ================================================================================================
262
263#[derive(Debug, Default, Clone, PartialEq, Eq)]
264struct BranchNode {
265 left: RpoDigest,
266 right: RpoDigest,
267}
268
269impl BranchNode {
270 fn parent(&self) -> RpoDigest {
271 Rpo256::merge(&[self.left, self.right])
272 }
273}