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