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