miden-crypto/src/merkle/merkle_tree.rs
src/merkle/merkle_tree.rsRaw
1use super::{Felt, MerkleError, 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 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 u32));
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 let two_nodes =
39 unsafe { slice::from_raw_parts(nodes.as_ptr() as *const [RpoDigest; 2], n) };
40
41 // calculate all internal tree nodes
42 for i in (1..n).rev() {
43 nodes[i] = Rpo256::merge(&two_nodes[i]).into();
44 }
45
46 Ok(Self { nodes })
47 }
48
49 // PUBLIC ACCESSORS
50 // --------------------------------------------------------------------------------------------
51
52 /// Returns the root of this Merkle tree.
53 pub fn root(&self) -> Word {
54 self.nodes[1]
55 }
56
57 /// Returns the depth of this Merkle tree.
58 ///
59 /// Merkle tree of depth 1 has two leaves, depth 2 has four leaves etc.
60 pub fn depth(&self) -> u32 {
61 log2(self.nodes.len() / 2)
62 }
63
64 /// Returns a node at the specified depth and index.
65 ///
66 /// # Errors
67 /// Returns an error if:
68 /// * The specified depth is greater than the depth of the tree.
69 /// * The specified index not valid for the specified depth.
70 pub fn get_node(&self, depth: u32, index: u64) -> Result<Word, MerkleError> {
71 if depth == 0 {
72 return Err(MerkleError::DepthTooSmall(depth));
73 } else if depth > self.depth() {
74 return Err(MerkleError::DepthTooBig(depth));
75 }
76 if index >= 2u64.pow(depth) {
77 return Err(MerkleError::InvalidIndex(depth, index));
78 }
79
80 let pos = 2_usize.pow(depth) + (index as usize);
81 Ok(self.nodes[pos])
82 }
83
84 /// Returns a Merkle path to the node at the specified depth and index. The node itself is
85 /// not included in the path.
86 ///
87 /// # Errors
88 /// Returns an error if:
89 /// * The specified depth is greater than the depth of the tree.
90 /// * The specified index not valid for the specified depth.
91 pub fn get_path(&self, depth: u32, index: u64) -> Result<Vec<Word>, MerkleError> {
92 if depth == 0 {
93 return Err(MerkleError::DepthTooSmall(depth));
94 } else if depth > self.depth() {
95 return Err(MerkleError::DepthTooBig(depth));
96 }
97 if index >= 2u64.pow(depth) {
98 return Err(MerkleError::InvalidIndex(depth, index));
99 }
100
101 let mut path = Vec::with_capacity(depth as usize);
102 let mut pos = 2_usize.pow(depth) + (index as usize);
103
104 while pos > 1 {
105 path.push(self.nodes[pos ^ 1]);
106 pos >>= 1;
107 }
108
109 Ok(path)
110 }
111
112 /// Replaces the leaf at the specified index with the provided value.
113 ///
114 /// # Errors
115 /// Returns an error if the specified index is not a valid leaf index for this tree.
116 pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<(), MerkleError> {
117 let depth = self.depth();
118 if index >= 2u64.pow(depth) {
119 return Err(MerkleError::InvalidIndex(depth, index));
120 }
121
122 let mut index = 2usize.pow(depth) + index as usize;
123 self.nodes[index] = value;
124
125 let n = self.nodes.len() / 2;
126 let two_nodes =
127 unsafe { slice::from_raw_parts(self.nodes.as_ptr() as *const [RpoDigest; 2], n) };
128
129 for _ in 0..depth {
130 index /= 2;
131 self.nodes[index] = Rpo256::merge(&two_nodes[index]).into();
132 }
133
134 Ok(())
135 }
136}
137
138// TESTS
139// ================================================================================================
140
141#[cfg(test)]
142mod tests {
143 use super::{
144 super::{int_to_node, Rpo256},
145 Word,
146 };
147
148 const LEAVES4: [Word; 4] = [
149 int_to_node(1),
150 int_to_node(2),
151 int_to_node(3),
152 int_to_node(4),
153 ];
154
155 const LEAVES8: [Word; 8] = [
156 int_to_node(1),
157 int_to_node(2),
158 int_to_node(3),
159 int_to_node(4),
160 int_to_node(5),
161 int_to_node(6),
162 int_to_node(7),
163 int_to_node(8),
164 ];
165
166 #[test]
167 fn build_merkle_tree() {
168 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
169 assert_eq!(8, tree.nodes.len());
170
171 // leaves were copied correctly
172 for (a, b) in tree.nodes.iter().skip(4).zip(LEAVES4.iter()) {
173 assert_eq!(a, b);
174 }
175
176 let (root, node2, node3) = compute_internal_nodes();
177
178 assert_eq!(root, tree.nodes[1]);
179 assert_eq!(node2, tree.nodes[2]);
180 assert_eq!(node3, tree.nodes[3]);
181
182 assert_eq!(root, tree.root());
183 }
184
185 #[test]
186 fn get_leaf() {
187 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
188
189 // check depth 2
190 assert_eq!(LEAVES4[0], tree.get_node(2, 0).unwrap());
191 assert_eq!(LEAVES4[1], tree.get_node(2, 1).unwrap());
192 assert_eq!(LEAVES4[2], tree.get_node(2, 2).unwrap());
193 assert_eq!(LEAVES4[3], tree.get_node(2, 3).unwrap());
194
195 // check depth 1
196 let (_, node2, node3) = compute_internal_nodes();
197
198 assert_eq!(node2, tree.get_node(1, 0).unwrap());
199 assert_eq!(node3, tree.get_node(1, 1).unwrap());
200 }
201
202 #[test]
203 fn get_path() {
204 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
205
206 let (_, node2, node3) = compute_internal_nodes();
207
208 // check depth 2
209 assert_eq!(vec![LEAVES4[1], node3], tree.get_path(2, 0).unwrap());
210 assert_eq!(vec![LEAVES4[0], node3], tree.get_path(2, 1).unwrap());
211 assert_eq!(vec![LEAVES4[3], node2], tree.get_path(2, 2).unwrap());
212 assert_eq!(vec![LEAVES4[2], node2], tree.get_path(2, 3).unwrap());
213
214 // check depth 1
215 assert_eq!(vec![node3], tree.get_path(1, 0).unwrap());
216 assert_eq!(vec![node2], tree.get_path(1, 1).unwrap());
217 }
218
219 #[test]
220 fn update_leaf() {
221 let mut tree = super::MerkleTree::new(LEAVES8.to_vec()).unwrap();
222
223 // update one leaf
224 let index = 3;
225 let new_node = int_to_node(9);
226 let mut expected_leaves = LEAVES8.to_vec();
227 expected_leaves[index as usize] = new_node;
228 let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap();
229
230 tree.update_leaf(index, new_node).unwrap();
231 assert_eq!(expected_tree.nodes, tree.nodes);
232
233 // update another leaf
234 let index = 6;
235 let new_node = int_to_node(10);
236 expected_leaves[index as usize] = new_node;
237 let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap();
238
239 tree.update_leaf(index, new_node).unwrap();
240 assert_eq!(expected_tree.nodes, tree.nodes);
241 }
242
243 // HELPER FUNCTIONS
244 // --------------------------------------------------------------------------------------------
245
246 fn compute_internal_nodes() -> (Word, Word, Word) {
247 let node2 = Rpo256::hash_elements(&[LEAVES4[0], LEAVES4[1]].concat());
248 let node3 = Rpo256::hash_elements(&[LEAVES4[2], LEAVES4[3]].concat());
249 let root = Rpo256::merge(&[node2, node3]);
250
251 (root.into(), node2.into(), node3.into())
252 }
253}