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 crate::{hash::Hasher, merkle::int_to_node, ElementHasher, HashFn, Word};
149
150 const LEAVES4: [Word; 4] = [
151 int_to_node(1),
152 int_to_node(2),
153 int_to_node(3),
154 int_to_node(4),
155 ];
156
157 const LEAVES8: [Word; 8] = [
158 int_to_node(1),
159 int_to_node(2),
160 int_to_node(3),
161 int_to_node(4),
162 int_to_node(5),
163 int_to_node(6),
164 int_to_node(7),
165 int_to_node(8),
166 ];
167
168 #[test]
169 fn build_merkle_tree() {
170 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
171 assert_eq!(8, tree.nodes.len());
172
173 // leaves were copied correctly
174 for (a, b) in tree.nodes.iter().skip(4).zip(LEAVES4.iter()) {
175 assert_eq!(a, b);
176 }
177
178 let (root, node2, node3) = compute_internal_nodes();
179
180 assert_eq!(root, tree.nodes[1]);
181 assert_eq!(node2, tree.nodes[2]);
182 assert_eq!(node3, tree.nodes[3]);
183
184 assert_eq!(root, tree.root());
185 }
186
187 #[test]
188 fn get_leaf() {
189 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
190
191 // check depth 2
192 assert_eq!(LEAVES4[0], tree.get_node(2, 0).unwrap());
193 assert_eq!(LEAVES4[1], tree.get_node(2, 1).unwrap());
194 assert_eq!(LEAVES4[2], tree.get_node(2, 2).unwrap());
195 assert_eq!(LEAVES4[3], tree.get_node(2, 3).unwrap());
196
197 // check depth 1
198 let (_, node2, node3) = compute_internal_nodes();
199
200 assert_eq!(node2, tree.get_node(1, 0).unwrap());
201 assert_eq!(node3, tree.get_node(1, 1).unwrap());
202 }
203
204 #[test]
205 fn get_path() {
206 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
207
208 let (_, node2, node3) = compute_internal_nodes();
209
210 // check depth 2
211 assert_eq!(vec![LEAVES4[1], node3], tree.get_path(2, 0).unwrap());
212 assert_eq!(vec![LEAVES4[0], node3], tree.get_path(2, 1).unwrap());
213 assert_eq!(vec![LEAVES4[3], node2], tree.get_path(2, 2).unwrap());
214 assert_eq!(vec![LEAVES4[2], node2], tree.get_path(2, 3).unwrap());
215
216 // check depth 1
217 assert_eq!(vec![node3], tree.get_path(1, 0).unwrap());
218 assert_eq!(vec![node2], tree.get_path(1, 1).unwrap());
219 }
220
221 #[test]
222 fn update_leaf() {
223 let mut tree = super::MerkleTree::new(LEAVES8.to_vec()).unwrap();
224
225 // update one leaf
226 let index = 3;
227 let new_node = int_to_node(9);
228 let mut expected_leaves = LEAVES8.to_vec();
229 expected_leaves[index as usize] = new_node;
230 let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap();
231
232 tree.update_leaf(index, new_node).unwrap();
233 assert_eq!(expected_tree.nodes, tree.nodes);
234
235 // update another leaf
236 let index = 6;
237 let new_node = int_to_node(10);
238 expected_leaves[index as usize] = new_node;
239 let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap();
240
241 tree.update_leaf(index, new_node).unwrap();
242 assert_eq!(expected_tree.nodes, tree.nodes);
243 }
244
245 // HELPER FUNCTIONS
246 // --------------------------------------------------------------------------------------------
247
248 fn compute_internal_nodes() -> (Word, Word, Word) {
249 let node2 = Hasher::hash_elements(&[LEAVES4[0], LEAVES4[1]].concat());
250 let node3 = Hasher::hash_elements(&[LEAVES4[2], LEAVES4[3]].concat());
251 let root = Hasher::merge(&[node2, node3]);
252
253 (root.into(), node2.into(), node3.into())
254 }
255}