miden-crypto/src/merkle/merkle_tree.rs
src/merkle/merkle_tree.rsRaw
1use super::{Felt, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word};
2use crate::{utils::uninit_vector, FieldElement};
3use core::slice;
4use winter_math::log2;
5
6// MERKLE TREE
7// ================================================================================================
8
9/// A fully-balanced binary Merkle tree (i.e., a tree where the number of leaves is a power of two).
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub struct MerkleTree {
12 pub(crate) nodes: Vec<Word>,
13}
14
15impl MerkleTree {
16 // CONSTRUCTOR
17 // --------------------------------------------------------------------------------------------
18 /// Returns a Merkle tree instantiated from the provided leaves.
19 ///
20 /// # Errors
21 /// Returns an error if the number of leaves is smaller than two or is not a power of two.
22 pub fn new(leaves: Vec<Word>) -> Result<Self, MerkleError> {
23 let n = leaves.len();
24 if n <= 1 {
25 return Err(MerkleError::DepthTooSmall(n as u8));
26 } else if !n.is_power_of_two() {
27 return Err(MerkleError::NumLeavesNotPowerOfTwo(n));
28 }
29
30 // create un-initialized vector to hold all tree nodes
31 let mut nodes = unsafe { uninit_vector(2 * n) };
32 nodes[0] = [Felt::ZERO; 4];
33
34 // copy leaves into the second part of the nodes vector
35 nodes[n..].copy_from_slice(&leaves);
36
37 // re-interpret nodes as an array of two nodes fused together
38 // Safety: `nodes` will never move here as it is not bound to an external lifetime (i.e.
39 // `self`).
40 let ptr = nodes.as_ptr() as *const [RpoDigest; 2];
41 let pairs = unsafe { slice::from_raw_parts(ptr, n) };
42
43 // calculate all internal tree nodes
44 for i in (1..n).rev() {
45 nodes[i] = Rpo256::merge(&pairs[i]).into();
46 }
47
48 Ok(Self { nodes })
49 }
50
51 // PUBLIC ACCESSORS
52 // --------------------------------------------------------------------------------------------
53
54 /// Returns the root of this Merkle tree.
55 pub fn root(&self) -> Word {
56 self.nodes[1]
57 }
58
59 /// Returns the depth of this Merkle tree.
60 ///
61 /// Merkle tree of depth 1 has two leaves, depth 2 has four leaves etc.
62 pub fn depth(&self) -> u8 {
63 log2(self.nodes.len() / 2) as u8
64 }
65
66 /// Returns a node at the specified depth and index value.
67 ///
68 /// # Errors
69 /// Returns an error if:
70 /// * The specified depth is greater than the depth of the tree.
71 /// * The specified index is not valid for the specified depth.
72 pub fn get_node(&self, index: NodeIndex) -> Result<Word, MerkleError> {
73 if index.is_root() {
74 return Err(MerkleError::DepthTooSmall(index.depth()));
75 } else if index.depth() > self.depth() {
76 return Err(MerkleError::DepthTooBig(index.depth() as u64));
77 } else if !index.is_valid() {
78 return Err(MerkleError::InvalidIndex(index));
79 }
80
81 let pos = index.to_scalar_index() as usize;
82 Ok(self.nodes[pos])
83 }
84
85 /// Returns a Merkle path to the node at the specified depth and index value. The node itself
86 /// is not included in the path.
87 ///
88 /// # Errors
89 /// Returns an error if:
90 /// * The specified depth is greater than the depth of the tree.
91 /// * The specified value is not valid for the specified depth.
92 pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
93 if index.is_root() {
94 return Err(MerkleError::DepthTooSmall(index.depth()));
95 } else if index.depth() > self.depth() {
96 return Err(MerkleError::DepthTooBig(index.depth() as u64));
97 } else if !index.is_valid() {
98 return Err(MerkleError::InvalidIndex(index));
99 }
100
101 // TODO should we create a helper in `NodeIndex` that will encapsulate traversal to root so
102 // we always use inlined `for` instead of `while`? the reason to use `for` is because its
103 // easier for the compiler to vectorize.
104 let mut path = Vec::with_capacity(index.depth() as usize);
105 for _ in 0..index.depth() {
106 let sibling = index.sibling().to_scalar_index() as usize;
107 path.push(self.nodes[sibling]);
108 index.move_up();
109 }
110
111 debug_assert!(
112 index.is_root(),
113 "the path walk must go all the way to the root"
114 );
115
116 Ok(path.into())
117 }
118
119 /// Replaces the leaf at the specified index with the provided value.
120 ///
121 /// # Errors
122 /// Returns an error if the specified index value is not a valid leaf value for this tree.
123 pub fn update_leaf<'a>(&'a mut self, index_value: u64, value: Word) -> Result<(), MerkleError> {
124 let depth = self.depth();
125 let mut index = NodeIndex::new(depth, index_value);
126 if !index.is_valid() {
127 return Err(MerkleError::InvalidIndex(index));
128 }
129
130 // we don't need to copy the pairs into a new address as we are logically guaranteed to not
131 // overlap write instructions. however, it's important to bind the lifetime of pairs to
132 // `self.nodes` so the compiler will never move one without moving the other.
133 debug_assert_eq!(self.nodes.len() & 1, 0);
134 let n = self.nodes.len() / 2;
135
136 // Safety: the length of nodes is guaranteed to contain pairs of words; hence, pairs of
137 // digests. we explicitly bind the lifetime here so we add an extra layer of guarantee that
138 // `self.nodes` will be moved only if `pairs` is moved as well. also, the algorithm is
139 // logically guaranteed to not overlap write positions as the write index is always half
140 // the index from which we read the digest input.
141 let ptr = self.nodes.as_ptr() as *const [RpoDigest; 2];
142 let pairs: &'a [[RpoDigest; 2]] = unsafe { slice::from_raw_parts(ptr, n) };
143
144 // update the current node
145 let pos = index.to_scalar_index() as usize;
146 self.nodes[pos] = value;
147
148 // traverse to the root, updating each node with the merged values of its parents
149 for _ in 0..index.depth() {
150 index.move_up();
151 let pos = index.to_scalar_index() as usize;
152 let value = Rpo256::merge(&pairs[pos]).into();
153 self.nodes[pos] = value;
154 }
155
156 Ok(())
157 }
158}
159
160// TESTS
161// ================================================================================================
162
163#[cfg(test)]
164mod tests {
165 use super::*;
166 use crate::merkle::int_to_node;
167 use core::mem::size_of;
168 use proptest::prelude::*;
169
170 const LEAVES4: [Word; 4] = [
171 int_to_node(1),
172 int_to_node(2),
173 int_to_node(3),
174 int_to_node(4),
175 ];
176
177 const LEAVES8: [Word; 8] = [
178 int_to_node(1),
179 int_to_node(2),
180 int_to_node(3),
181 int_to_node(4),
182 int_to_node(5),
183 int_to_node(6),
184 int_to_node(7),
185 int_to_node(8),
186 ];
187
188 #[test]
189 fn build_merkle_tree() {
190 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
191 assert_eq!(8, tree.nodes.len());
192
193 // leaves were copied correctly
194 for (a, b) in tree.nodes.iter().skip(4).zip(LEAVES4.iter()) {
195 assert_eq!(a, b);
196 }
197
198 let (root, node2, node3) = compute_internal_nodes();
199
200 assert_eq!(root, tree.nodes[1]);
201 assert_eq!(node2, tree.nodes[2]);
202 assert_eq!(node3, tree.nodes[3]);
203
204 assert_eq!(root, tree.root());
205 }
206
207 #[test]
208 fn get_leaf() {
209 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
210
211 // check depth 2
212 assert_eq!(LEAVES4[0], tree.get_node(NodeIndex::new(2, 0)).unwrap());
213 assert_eq!(LEAVES4[1], tree.get_node(NodeIndex::new(2, 1)).unwrap());
214 assert_eq!(LEAVES4[2], tree.get_node(NodeIndex::new(2, 2)).unwrap());
215 assert_eq!(LEAVES4[3], tree.get_node(NodeIndex::new(2, 3)).unwrap());
216
217 // check depth 1
218 let (_, node2, node3) = compute_internal_nodes();
219
220 assert_eq!(node2, tree.get_node(NodeIndex::new(1, 0)).unwrap());
221 assert_eq!(node3, tree.get_node(NodeIndex::new(1, 1)).unwrap());
222 }
223
224 #[test]
225 fn get_path() {
226 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
227
228 let (_, node2, node3) = compute_internal_nodes();
229
230 // check depth 2
231 assert_eq!(
232 vec![LEAVES4[1], node3],
233 *tree.get_path(NodeIndex::new(2, 0)).unwrap()
234 );
235 assert_eq!(
236 vec![LEAVES4[0], node3],
237 *tree.get_path(NodeIndex::new(2, 1)).unwrap()
238 );
239 assert_eq!(
240 vec![LEAVES4[3], node2],
241 *tree.get_path(NodeIndex::new(2, 2)).unwrap()
242 );
243 assert_eq!(
244 vec![LEAVES4[2], node2],
245 *tree.get_path(NodeIndex::new(2, 3)).unwrap()
246 );
247
248 // check depth 1
249 assert_eq!(vec![node3], *tree.get_path(NodeIndex::new(1, 0)).unwrap());
250 assert_eq!(vec![node2], *tree.get_path(NodeIndex::new(1, 1)).unwrap());
251 }
252
253 #[test]
254 fn update_leaf() {
255 let mut tree = super::MerkleTree::new(LEAVES8.to_vec()).unwrap();
256
257 // update one leaf
258 let value = 3;
259 let new_node = int_to_node(9);
260 let mut expected_leaves = LEAVES8.to_vec();
261 expected_leaves[value as usize] = new_node;
262 let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap();
263
264 tree.update_leaf(value, new_node).unwrap();
265 assert_eq!(expected_tree.nodes, tree.nodes);
266
267 // update another leaf
268 let value = 6;
269 let new_node = int_to_node(10);
270 expected_leaves[value as usize] = new_node;
271 let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap();
272
273 tree.update_leaf(value, new_node).unwrap();
274 assert_eq!(expected_tree.nodes, tree.nodes);
275 }
276
277 proptest! {
278 #[test]
279 fn arbitrary_word_can_be_represented_as_digest(
280 a in prop::num::u64::ANY,
281 b in prop::num::u64::ANY,
282 c in prop::num::u64::ANY,
283 d in prop::num::u64::ANY,
284 ) {
285 // this test will assert the memory equivalence between word and digest.
286 // it is used to safeguard the `[MerkleTee::update_leaf]` implementation
287 // that assumes this equivalence.
288
289 // build a word and copy it to another address as digest
290 let word = [Felt::new(a), Felt::new(b), Felt::new(c), Felt::new(d)];
291 let digest = RpoDigest::from(word);
292
293 // assert the addresses are different
294 let word_ptr = (&word).as_ptr() as *const u8;
295 let digest_ptr = (&digest).as_ptr() as *const u8;
296 assert_ne!(word_ptr, digest_ptr);
297
298 // compare the bytes representation
299 let word_bytes = unsafe { slice::from_raw_parts(word_ptr, size_of::<Word>()) };
300 let digest_bytes = unsafe { slice::from_raw_parts(digest_ptr, size_of::<RpoDigest>()) };
301 assert_eq!(word_bytes, digest_bytes);
302 }
303 }
304
305 // HELPER FUNCTIONS
306 // --------------------------------------------------------------------------------------------
307
308 fn compute_internal_nodes() -> (Word, Word, Word) {
309 let node2 = Rpo256::hash_elements(&[LEAVES4[0], LEAVES4[1]].concat());
310 let node3 = Rpo256::hash_elements(&[LEAVES4[2], LEAVES4[3]].concat());
311 let root = Rpo256::merge(&[node2, node3]);
312
313 (root.into(), node2.into(), node3.into())
314 }
315}