miden-crypto/src/merkle/simple_smt/mod.rs
src/merkle/simple_smt/mod.rsRaw
1use super::{
2 BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256,
3 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/// The root of the tree is recomputed on each new leaf update.
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct SimpleSmt {
16 depth: u8,
17 root: Word,
18 leaves: BTreeMap<u64, Word>,
19 branches: BTreeMap<NodeIndex, BranchNode>,
20 empty_hashes: Vec<RpoDigest>,
21}
22
23#[derive(Debug, Default, Clone, PartialEq, Eq)]
24struct BranchNode {
25 left: RpoDigest,
26 right: RpoDigest,
27}
28
29impl BranchNode {
30 fn parent(&self) -> RpoDigest {
31 Rpo256::merge(&[self.left, self.right])
32 }
33}
34
35impl SimpleSmt {
36 // CONSTANTS
37 // --------------------------------------------------------------------------------------------
38
39 /// Minimum supported depth.
40 pub const MIN_DEPTH: u8 = 1;
41
42 /// Maximum supported depth.
43 pub const MAX_DEPTH: u8 = 64;
44
45 // CONSTRUCTORS
46 // --------------------------------------------------------------------------------------------
47
48 /// Creates a new simple SMT with the provided depth.
49 pub fn new(depth: u8) -> Result<Self, MerkleError> {
50 // validate the range of the depth.
51 if depth < Self::MIN_DEPTH {
52 return Err(MerkleError::DepthTooSmall(depth));
53 } else if Self::MAX_DEPTH < depth {
54 return Err(MerkleError::DepthTooBig(depth as u64));
55 }
56
57 let empty_hashes = EmptySubtreeRoots::empty_hashes(depth).to_vec();
58 let root = empty_hashes[0].into();
59
60 Ok(Self {
61 root,
62 depth,
63 empty_hashes,
64 leaves: BTreeMap::new(),
65 branches: BTreeMap::new(),
66 })
67 }
68
69 /// Appends the provided entries as leaves of the tree.
70 ///
71 /// # Errors
72 ///
73 /// The function will fail if the provided entries count exceed the maximum tree capacity, that
74 /// is `2^{depth}`.
75 pub fn with_leaves<R, I>(mut self, entries: R) -> Result<Self, MerkleError>
76 where
77 R: IntoIterator<IntoIter = I>,
78 I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
79 {
80 // check if the leaves count will fit the depth setup
81 let mut entries = entries.into_iter();
82 let max = 1 << self.depth.min(63);
83 if entries.len() > max {
84 return Err(MerkleError::InvalidEntriesCount(max, entries.len()));
85 }
86
87 // append leaves and return
88 entries.try_for_each(|(key, leaf)| self.insert_leaf(key, leaf))?;
89 Ok(self)
90 }
91
92 /// Replaces the internal empty digests used when a given depth doesn't contain a node.
93 pub fn with_empty_subtrees<I>(mut self, hashes: I) -> Self
94 where
95 I: IntoIterator<Item = RpoDigest>,
96 {
97 self.replace_empty_subtrees(hashes.into_iter().collect());
98 self
99 }
100
101 // PUBLIC ACCESSORS
102 // --------------------------------------------------------------------------------------------
103
104 /// Returns the root of this Merkle tree.
105 pub const fn root(&self) -> Word {
106 self.root
107 }
108
109 /// Returns the depth of this Merkle tree.
110 pub const fn depth(&self) -> u8 {
111 self.depth
112 }
113
114 // PROVIDERS
115 // --------------------------------------------------------------------------------------------
116
117 /// Returns the set count of the keys of the leaves.
118 pub fn leaves_count(&self) -> usize {
119 self.leaves.len()
120 }
121
122 /// Returns a node at the specified key
123 ///
124 /// # Errors
125 /// Returns an error if:
126 /// * The specified depth is greater than the depth of the tree.
127 pub fn get_node(&self, index: &NodeIndex) -> Result<Word, 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 self.get_leaf_node(index.value())
134 .or_else(|| {
135 self.empty_hashes
136 .get(index.depth() as usize)
137 .copied()
138 .map(Word::from)
139 })
140 .ok_or(MerkleError::InvalidIndex(*index))
141 } else {
142 let branch_node = self.get_branch_node(index);
143 Ok(Rpo256::merge(&[branch_node.left, branch_node.right]).into())
144 }
145 }
146
147 /// Returns a Merkle path from the node at the specified key to the root. The node itself is
148 /// not included in the path.
149 ///
150 /// # Errors
151 /// Returns an error if:
152 /// * The specified depth is greater than the depth of the tree.
153 pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
154 if index.is_root() {
155 return Err(MerkleError::DepthTooSmall(index.depth()));
156 } else if index.depth() > self.depth() {
157 return Err(MerkleError::DepthTooBig(index.depth() as u64));
158 }
159
160 let mut path = Vec::with_capacity(index.depth() as usize);
161 for _ in 0..index.depth() {
162 let is_right = index.is_value_odd();
163 index.move_up();
164 let BranchNode { left, right } = self.get_branch_node(&index);
165 let value = if is_right { left } else { right };
166 path.push(*value);
167 }
168 Ok(path.into())
169 }
170
171 /// Return a Merkle path from the leaf at the specified key to the root. The leaf itself is not
172 /// included in the path.
173 ///
174 /// # Errors
175 /// Returns an error if:
176 /// * The specified key does not exist as a leaf node.
177 pub fn get_leaf_path(&self, key: u64) -> Result<MerklePath, MerkleError> {
178 self.get_path(NodeIndex::new(self.depth(), key))
179 }
180
181 /// Iterator over the inner nodes of the [SimpleSmt].
182 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
183 self.branches.values().map(|e| InnerNodeInfo {
184 value: e.parent().into(),
185 left: e.left.into(),
186 right: e.right.into(),
187 })
188 }
189
190 // STATE MUTATORS
191 // --------------------------------------------------------------------------------------------
192
193 /// Replaces the leaf located at the specified key, and recomputes hashes by walking up the tree
194 ///
195 /// # Errors
196 /// Returns an error if the specified key is not a valid leaf index for this tree.
197 pub fn update_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> {
198 if !self.check_leaf_node_exists(key) {
199 return Err(MerkleError::InvalidIndex(NodeIndex::new(self.depth(), key)));
200 }
201 self.insert_leaf(key, value)?;
202
203 Ok(())
204 }
205
206 /// Inserts a leaf located at the specified key, and recomputes hashes by walking up the tree
207 pub fn insert_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> {
208 self.insert_leaf_node(key, value);
209
210 // TODO consider using a map `index |-> word` instead of `index |-> (word, word)`
211 let mut index = NodeIndex::new(self.depth(), key);
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 {
218 (left, value)
219 } else {
220 (value, right)
221 };
222 self.insert_branch_node(index, left, right);
223 value = Rpo256::merge(&[left, right]);
224 }
225 self.root = value.into();
226 Ok(())
227 }
228
229 // HELPER METHODS
230 // --------------------------------------------------------------------------------------------
231
232 fn replace_empty_subtrees(&mut self, hashes: Vec<RpoDigest>) {
233 self.empty_hashes = hashes;
234 }
235
236 fn check_leaf_node_exists(&self, key: u64) -> bool {
237 self.leaves.contains_key(&key)
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) {
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}