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 /// Returns an iterator over the inner nodes of this Merkle tree.
186 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
187 self.branches.values().map(|e| InnerNodeInfo {
188 value: e.parent().into(),
189 left: e.left.into(),
190 right: e.right.into(),
191 })
192 }
193
194 // STATE MUTATORS
195 // --------------------------------------------------------------------------------------------
196
197 /// Updates value of the leaf at the specified index returning the old leaf value.
198 ///
199 /// This also recomputes all hashes between the leaf and the root, updating the root itself.
200 ///
201 /// # Errors
202 /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
203 pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<Word, MerkleError> {
204 let old_value = self.insert_leaf_node(index, value).unwrap_or(EMPTY_WORD);
205
206 // if the old value and new value are the same, there is nothing to update
207 if value == old_value {
208 return Ok(value);
209 }
210
211 let mut index = NodeIndex::new(self.depth(), index)?;
212 let mut value = RpoDigest::from(value);
213 for _ in 0..index.depth() {
214 let is_right = index.is_value_odd();
215 index.move_up();
216 let BranchNode { left, right } = self.get_branch_node(&index);
217 let (left, right) = if is_right { (left, value) } else { (value, right) };
218 self.insert_branch_node(index, left, right);
219 value = Rpo256::merge(&[left, right]);
220 }
221 self.root = value.into();
222 Ok(old_value)
223 }
224
225 // HELPER METHODS
226 // --------------------------------------------------------------------------------------------
227
228 fn get_leaf_node(&self, key: u64) -> Option<Word> {
229 self.leaves.get(&key).copied()
230 }
231
232 fn insert_leaf_node(&mut self, key: u64, node: Word) -> Option<Word> {
233 self.leaves.insert(key, node)
234 }
235
236 fn get_branch_node(&self, index: &NodeIndex) -> BranchNode {
237 self.branches.get(index).cloned().unwrap_or_else(|| {
238 let node = self.empty_hashes[index.depth() as usize + 1];
239 BranchNode {
240 left: node,
241 right: node,
242 }
243 })
244 }
245
246 fn insert_branch_node(&mut self, index: NodeIndex, left: RpoDigest, right: RpoDigest) {
247 let branch = BranchNode { left, right };
248 self.branches.insert(index, branch);
249 }
250}
251
252// BRANCH NODE
253// ================================================================================================
254
255#[derive(Debug, Default, Clone, PartialEq, Eq)]
256struct BranchNode {
257 left: RpoDigest,
258 right: RpoDigest,
259}
260
261impl BranchNode {
262 fn parent(&self) -> RpoDigest {
263 Rpo256::merge(&[self.left, self.right])
264 }
265}