miden-crypto/src/merkle/partial_mt/mod.rs
src/merkle/partial_mt/mod.rsRaw
1use super::{
2 BTreeMap, BTreeSet, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec,
3 Word, EMPTY_WORD,
4};
5
6#[cfg(test)]
7mod tests;
8
9// PARTIAL MERKLE TREE
10// ================================================================================================
11
12/// A partial Merkle tree with NodeIndex keys and 4-element RpoDigest leaf values.
13///
14/// The root of the tree is recomputed on each new leaf update.
15pub struct PartialMerkleTree {
16 root: RpoDigest,
17 max_depth: u8,
18 nodes: BTreeMap<NodeIndex, RpoDigest>,
19 leaves: BTreeSet<NodeIndex>,
20}
21
22impl Default for PartialMerkleTree {
23 fn default() -> Self {
24 Self::new()
25 }
26}
27
28impl PartialMerkleTree {
29 // CONSTANTS
30 // --------------------------------------------------------------------------------------------
31
32 /// An RpoDigest consisting of 4 ZERO elements.
33 pub const EMPTY_DIGEST: RpoDigest = RpoDigest::new(EMPTY_WORD);
34
35 /// Minimum supported depth.
36 pub const MIN_DEPTH: u8 = 1;
37
38 /// Maximum supported depth.
39 pub const MAX_DEPTH: u8 = 64;
40
41 // CONSTRUCTORS
42 // --------------------------------------------------------------------------------------------
43
44 /// Returns a new emply [PartialMerkleTree].
45 pub fn new() -> Self {
46 PartialMerkleTree {
47 root: Self::EMPTY_DIGEST,
48 max_depth: 0,
49 nodes: BTreeMap::new(),
50 leaves: BTreeSet::new(),
51 }
52 }
53
54 /// Returns a new [PartialMerkleTree] instantiated with leaves set as specified by the provided
55 /// entries.
56 ///
57 /// # Errors
58 /// Returns an error if:
59 /// - If the depth is 0 or is greater than 64.
60 /// - The number of entries exceeds the maximum tree capacity, that is 2^{depth}.
61 /// - The provided entries contain multiple values for the same key.
62 pub fn with_leaves<R, I>(entries: R) -> Result<Self, MerkleError>
63 where
64 R: IntoIterator<IntoIter = I>,
65 I: Iterator<Item = (NodeIndex, RpoDigest)> + ExactSizeIterator,
66 {
67 // create an empty tree
68 let mut tree = PartialMerkleTree::new();
69
70 // check if the number of leaves can be accommodated by the tree's depth; we use a min
71 // depth of 63 because we consider passing in a vector of size 2^64 infeasible.
72 let entries = entries.into_iter();
73 let max = (1_u64 << 63) as usize;
74 if entries.len() > max {
75 return Err(MerkleError::InvalidNumEntries(max, entries.len()));
76 }
77
78 for (node_index, rpo_digest) in entries {
79 let old_value = tree.update_leaf(node_index, rpo_digest)?;
80 if old_value != Self::EMPTY_DIGEST {
81 return Err(MerkleError::DuplicateValuesForIndex(node_index.value()));
82 }
83 }
84 Ok(tree)
85 }
86
87 // PUBLIC ACCESSORS
88 // --------------------------------------------------------------------------------------------
89
90 /// Returns the root of this Merkle tree.
91 pub fn root(&self) -> Word {
92 self.root.into()
93 }
94
95 /// Returns the depth of this Merkle tree.
96 // TODO: maybe it's better to rename it to the `max_depth`
97 pub fn depth(&self) -> u8 {
98 self.max_depth
99 }
100
101 /// Returns a node at the specified NodeIndex.
102 ///
103 /// # Errors
104 /// Returns an error if the specified NodeIndex is not contained in the nodes map.
105 pub fn get_node(&self, index: NodeIndex) -> Result<Word, MerkleError> {
106 self.nodes
107 .get(&index)
108 .ok_or(MerkleError::NodeNotInSet(index))
109 .map(|hash| **hash)
110 }
111
112 /// Returns a value of the leaf at the specified NodeIndex.
113 ///
114 /// # Errors
115 /// Returns an error if the NodeIndex is not contained in the leaves set.
116 pub fn get_leaf(&self, index: NodeIndex) -> Result<Word, MerkleError> {
117 if !self.leaves.contains(&index) {
118 // This error not really suitable in this situation, should I create a new error?
119 Err(MerkleError::InvalidIndex {
120 depth: index.depth(),
121 value: index.value(),
122 })
123 } else {
124 self.nodes
125 .get(&index)
126 .ok_or(MerkleError::NodeNotInSet(index))
127 .map(|hash| **hash)
128 }
129 }
130
131 /// Returns a map of the all
132 pub fn paths(&self) -> Result<BTreeMap<&NodeIndex, MerklePath>, MerkleError> {
133 let mut paths = BTreeMap::new();
134 for leaf_index in self.leaves.iter() {
135 let index = *leaf_index;
136 paths.insert(leaf_index, self.get_path(index)?);
137 }
138 Ok(paths)
139 }
140
141 /// Returns a Merkle path from the node at the specified index to the root.
142 ///
143 /// The node itself is not included in the path.
144 ///
145 /// # Errors
146 /// Returns an error if:
147 /// - the specified index has depth set to 0 or the depth is greater than the depth of this
148 /// Merkle tree.
149 /// - the specified index is not contained in the nodes map.
150 pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
151 if index.is_root() {
152 return Err(MerkleError::DepthTooSmall(index.depth()));
153 } else if index.depth() > self.depth() {
154 return Err(MerkleError::DepthTooBig(index.depth() as u64));
155 }
156
157 if !self.nodes.contains_key(&index) {
158 return Err(MerkleError::NodeNotInSet(index));
159 }
160
161 let mut path = Vec::new();
162 for _ in 0..index.depth() {
163 let is_right = index.is_value_odd();
164 let sibling_index = if is_right {
165 NodeIndex::new(index.depth(), index.value() - 1)?
166 } else {
167 NodeIndex::new(index.depth(), index.value() + 1)?
168 };
169 index.move_up();
170 let sibling_hash =
171 self.nodes.get(&sibling_index).cloned().unwrap_or(Self::EMPTY_DIGEST);
172 path.push(Word::from(sibling_hash));
173 }
174 Ok(MerklePath::new(path))
175 }
176
177 // ITERATORS
178 // --------------------------------------------------------------------------------------------
179
180 /// Returns an iterator over the leaves of this [PartialMerkleTree].
181 pub fn leaves(&self) -> impl Iterator<Item = (NodeIndex, &Word)> {
182 self.nodes
183 .iter()
184 .filter(|(index, _)| self.leaves.contains(index))
185 .map(|(index, hash)| (*index, &(**hash)))
186 }
187
188 /// Returns an iterator over the inner nodes of this Merkle tree.
189 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
190 let inner_nodes = self.nodes.iter().filter(|(index, _)| !self.leaves.contains(index));
191 inner_nodes.map(|(index, digest)| {
192 let left_index = NodeIndex::new(index.depth() + 1, index.value() * 2)
193 .expect("Failure to get left child index");
194 let right_index = NodeIndex::new(index.depth() + 1, index.value() * 2 + 1)
195 .expect("Failure to get right child index");
196 let left_hash = self.nodes.get(&left_index).cloned().unwrap_or(Self::EMPTY_DIGEST);
197 let right_hash = self.nodes.get(&right_index).cloned().unwrap_or(Self::EMPTY_DIGEST);
198 InnerNodeInfo {
199 value: **digest,
200 left: *left_hash,
201 right: *right_hash,
202 }
203 })
204 }
205
206 // STATE MUTATORS
207 // --------------------------------------------------------------------------------------------
208
209 /// Updates value of the leaf at the specified index returning the old leaf value.
210 ///
211 /// This also recomputes all hashes between the leaf and the root, updating the root itself.
212 pub fn update_leaf(
213 &mut self,
214 node_index: NodeIndex,
215 value: RpoDigest,
216 ) -> Result<RpoDigest, MerkleError> {
217 // check correctness of the depth and update it
218 Self::check_depth(node_index.depth())?;
219 self.update_depth(node_index.depth());
220
221 // insert NodeIndex to the leaves Set
222 self.leaves.insert(node_index);
223
224 // add node value to the nodes Map
225 let old_value = self.nodes.insert(node_index, value).unwrap_or(Self::EMPTY_DIGEST);
226
227 // if the old value and new value are the same, there is nothing to update
228 if value == old_value {
229 return Ok(value);
230 }
231
232 let mut node_index = node_index;
233 let mut value = value;
234 for _ in 0..node_index.depth() {
235 let is_right = node_index.is_value_odd();
236 let (left, right) = if is_right {
237 let left_index = NodeIndex::new(node_index.depth(), node_index.value() - 1)?;
238 (self.nodes.get(&left_index).cloned().unwrap_or(Self::EMPTY_DIGEST), value)
239 } else {
240 let right_index = NodeIndex::new(node_index.depth(), node_index.value() + 1)?;
241 (value, self.nodes.get(&right_index).cloned().unwrap_or(Self::EMPTY_DIGEST))
242 };
243 node_index.move_up();
244 value = Rpo256::merge(&[left, right]);
245 self.nodes.insert(node_index, value);
246 }
247
248 self.root = value;
249 Ok(old_value)
250 }
251
252 // HELPER METHODS
253 // --------------------------------------------------------------------------------------------
254
255 /// Updates depth value with the maximum of current and provided depth.
256 fn update_depth(&mut self, new_depth: u8) {
257 self.max_depth = new_depth.max(self.max_depth);
258 }
259
260 /// Returns an error if the depth is 0 or is greater than 64.
261 fn check_depth(depth: u8) -> Result<(), MerkleError> {
262 // validate the range of the depth.
263 if depth < Self::MIN_DEPTH {
264 return Err(MerkleError::DepthTooSmall(depth));
265 } else if Self::MAX_DEPTH < depth {
266 return Err(MerkleError::DepthTooBig(depth as u64));
267 }
268 Ok(())
269 }
270}