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