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