miden-crypto/crypto/src/merkle/merkle_tree.rs
crypto/src/merkle/merkle_tree.rsRaw
1use super::MerkleError;
2use crate::{
3 hash::{merge, Digest},
4 Felt, FieldElement, Word,
5};
6use core::slice;
7use winter_utils::uninit_vector;
8use winterfell::math::log2;
9
10// MERKLE TREE
11// ================================================================================================
12
13/// A fully-balanced binary Merkle tree (i.e., a tree where the number of leaves is a power of two).
14///
15/// This struct is intended to be used as one of the variants of the MerkleSet enum.
16#[derive(Clone, Debug)]
17pub struct MerkleTree {
18 nodes: Vec<Word>,
19}
20
21impl MerkleTree {
22 // CONSTRUCTOR
23 // --------------------------------------------------------------------------------------------
24 /// Returns a Merkle tree instantiated from the provided leaves.
25 ///
26 /// # Errors
27 /// Returns an error if the number of leaves is smaller than two or is not a power of two.
28 pub fn new(leaves: Vec<Word>) -> Result<Self, MerkleError> {
29 let n = leaves.len();
30 if n <= 1 {
31 return Err(MerkleError::DepthTooSmall);
32 } else if !n.is_power_of_two() {
33 return Err(MerkleError::NumLeavesNotPowerOfTwo(n));
34 }
35
36 // create un-initialized vector to hold all tree nodes
37 let mut nodes = unsafe { uninit_vector(2 * n) };
38 nodes[0] = [Felt::ZERO; 4];
39
40 // copy leaves into the second part of the nodes vector
41 nodes[n..].copy_from_slice(&leaves);
42
43 // re-interpret nodes as an array of two nodes fused together
44 let two_nodes = unsafe { slice::from_raw_parts(nodes.as_ptr() as *const [Digest; 2], n) };
45
46 // calculate all internal tree nodes
47 for i in (1..n).rev() {
48 nodes[i] = merge(&two_nodes[i]).into();
49 }
50
51 Ok(Self { nodes })
52 }
53
54 // PUBLIC ACCESSORS
55 // --------------------------------------------------------------------------------------------
56
57 /// Returns the root of this Merkle tree.
58 pub fn root(&self) -> Word {
59 self.nodes[1]
60 }
61
62 /// Returns the depth of this Merkle tree.
63 ///
64 /// Merkle tree of depth 1 has two leaves, depth 2 has four leaves etc.
65 pub fn depth(&self) -> u32 {
66 log2(self.nodes.len() / 2)
67 }
68
69 /// Returns a node at the specified depth and index.
70 ///
71 /// # Errors
72 /// Returns an error if:
73 /// * The specified depth is greater than the depth of the tree.
74 /// * The specified index not valid for the specified depth.
75 pub fn get_node(&self, depth: u32, index: u64) -> Result<Word, MerkleError> {
76 if depth == 0 {
77 return Err(MerkleError::DepthTooSmall);
78 } else if depth > self.depth() {
79 return Err(MerkleError::DepthTooBig(depth));
80 }
81 if index >= 2u64.pow(depth) {
82 return Err(MerkleError::InvalidIndex(depth, index));
83 }
84
85 let pos = 2usize.pow(depth as u32) + (index as usize);
86 Ok(self.nodes[pos])
87 }
88
89 /// Returns a Merkle path to the node at the specified depth and index. The node itself is
90 /// not included in the path.
91 ///
92 /// # Errors
93 /// Returns an error if:
94 /// * The specified depth is greater than the depth of the tree.
95 /// * The specified index not valid for the specified depth.
96 pub fn get_path(&self, depth: u32, index: u64) -> Result<Vec<Word>, MerkleError> {
97 if depth == 0 {
98 return Err(MerkleError::DepthTooSmall);
99 } else if depth > self.depth() {
100 return Err(MerkleError::DepthTooBig(depth));
101 }
102 if index >= 2u64.pow(depth) {
103 return Err(MerkleError::InvalidIndex(depth, index));
104 }
105
106 let mut path = Vec::with_capacity(depth as usize);
107 let mut pos = 2usize.pow(depth as u32) + (index as usize);
108
109 while pos > 1 {
110 path.push(self.nodes[pos ^ 1]);
111 pos >>= 1;
112 }
113
114 Ok(path)
115 }
116
117 /// Replaces the leaf at the specified index with the provided value.
118 ///
119 /// # Errors
120 /// Returns an error if the specified index is not a valid leaf index for this tree.
121 pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<(), MerkleError> {
122 let depth = self.depth();
123 if index >= 2u64.pow(depth) {
124 return Err(MerkleError::InvalidIndex(depth, index));
125 }
126
127 let mut index = 2usize.pow(depth) + index as usize;
128 self.nodes[index] = value;
129
130 let n = self.nodes.len() / 2;
131 let two_nodes =
132 unsafe { slice::from_raw_parts(self.nodes.as_ptr() as *const [Digest; 2], n) };
133
134 for _ in 0..depth {
135 index /= 2;
136 self.nodes[index] = merge(&two_nodes[index]).into();
137 }
138
139 Ok(())
140 }
141}
142
143// TESTS
144// ================================================================================================
145
146#[cfg(test)]
147mod tests {
148 use super::{Felt, FieldElement, Word};
149 use crate::hash::{ElementHasher, HashFn, Hasher};
150
151 const LEAVES4: [Word; 4] = [
152 int_to_node(1),
153 int_to_node(2),
154 int_to_node(3),
155 int_to_node(4),
156 ];
157
158 const LEAVES8: [Word; 8] = [
159 int_to_node(1),
160 int_to_node(2),
161 int_to_node(3),
162 int_to_node(4),
163 int_to_node(5),
164 int_to_node(6),
165 int_to_node(7),
166 int_to_node(8),
167 ];
168
169 #[test]
170 fn build_merkle_tree() {
171 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
172 assert_eq!(8, tree.nodes.len());
173
174 // leaves were copied correctly
175 for (a, b) in tree.nodes.iter().skip(4).zip(LEAVES4.iter()) {
176 assert_eq!(a, b);
177 }
178
179 let (root, node2, node3) = compute_internal_nodes();
180
181 assert_eq!(root, tree.nodes[1]);
182 assert_eq!(node2, tree.nodes[2]);
183 assert_eq!(node3, tree.nodes[3]);
184
185 assert_eq!(root, tree.root());
186 }
187
188 #[test]
189 fn get_leaf() {
190 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
191
192 // check depth 2
193 assert_eq!(LEAVES4[0], tree.get_node(2, 0).unwrap());
194 assert_eq!(LEAVES4[1], tree.get_node(2, 1).unwrap());
195 assert_eq!(LEAVES4[2], tree.get_node(2, 2).unwrap());
196 assert_eq!(LEAVES4[3], tree.get_node(2, 3).unwrap());
197
198 // check depth 1
199 let (_, node2, node3) = compute_internal_nodes();
200
201 assert_eq!(node2, tree.get_node(1, 0).unwrap());
202 assert_eq!(node3, tree.get_node(1, 1).unwrap());
203 }
204
205 #[test]
206 fn get_path() {
207 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
208
209 let (_, node2, node3) = compute_internal_nodes();
210
211 // check depth 2
212 assert_eq!(vec![LEAVES4[1], node3], tree.get_path(2, 0).unwrap());
213 assert_eq!(vec![LEAVES4[0], node3], tree.get_path(2, 1).unwrap());
214 assert_eq!(vec![LEAVES4[3], node2], tree.get_path(2, 2).unwrap());
215 assert_eq!(vec![LEAVES4[2], node2], tree.get_path(2, 3).unwrap());
216
217 // check depth 1
218 assert_eq!(vec![node3], tree.get_path(1, 0).unwrap());
219 assert_eq!(vec![node2], tree.get_path(1, 1).unwrap());
220 }
221
222 #[test]
223 fn update_leaf() {
224 let mut tree = super::MerkleTree::new(LEAVES8.to_vec()).unwrap();
225
226 // update one leaf
227 let index = 3;
228 let new_node = int_to_node(9);
229 let mut expected_leaves = LEAVES8.to_vec();
230 expected_leaves[index as usize] = new_node;
231 let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap();
232
233 tree.update_leaf(index, new_node).unwrap();
234 assert_eq!(expected_tree.nodes, tree.nodes);
235
236 // update another leaf
237 let index = 6;
238 let new_node = int_to_node(10);
239 expected_leaves[index as usize] = new_node;
240 let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap();
241
242 tree.update_leaf(index, new_node).unwrap();
243 assert_eq!(expected_tree.nodes, tree.nodes);
244 }
245
246 // HELPER FUNCTIONS
247 // --------------------------------------------------------------------------------------------
248
249 fn compute_internal_nodes() -> (Word, Word, Word) {
250 let node2 = Hasher::hash_elements(&[LEAVES4[0], LEAVES4[1]].concat());
251 let node3 = Hasher::hash_elements(&[LEAVES4[2], LEAVES4[3]].concat());
252 let root = Hasher::merge(&[node2, node3]);
253
254 (root.into(), node2.into(), node3.into())
255 }
256
257 const fn int_to_node(value: u64) -> Word {
258 [Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO]
259 }
260}