fri-commitment/src/merkletree.rs
src/merkletree.rsRaw
1// merkletree.rs implements a simple binary insert-only merkletree in which the leafs positions is
2// determined by the leaf value binary representation. Inspired by
3// https://docs.iden3.io/publications/pdfs/Merkle-Tree.pdf (which can be found implemented in
4// https://github.com/vocdoni/arbo).
5
6use ark_ff::{BigInteger, PrimeField};
7use ark_serialize::CanonicalSerialize;
8use ark_std::log2;
9use ark_std::marker::PhantomData;
10use sha3::{Digest, Keccak256};
11
12// abstraction to set the hash function used
13pub trait Hash<F>: Clone {
14 fn hash(_in: &[F]) -> F;
15}
16
17#[derive(Clone, Copy, Debug)]
18pub struct Keccak256Hash<F: PrimeField> {
19 phantom: PhantomData<F>,
20}
21impl<F: PrimeField> Hash<F> for Keccak256Hash<F> {
22 fn hash(_in: &[F]) -> F {
23 let mut buf = vec![];
24 _in.serialize_uncompressed(&mut buf).unwrap();
25
26 let mut h = Keccak256::default();
27 h.update(buf);
28 let r = h.finalize();
29 let out = F::from_le_bytes_mod_order(&r);
30 out
31 }
32}
33
34#[derive(Clone, Debug)]
35pub struct Node<F: PrimeField, H: Hash<F>> {
36 phantom: PhantomData<H>,
37 hash: F,
38 left: Option<Box<Node<F, H>>>,
39 right: Option<Box<Node<F, H>>>,
40 value: Option<F>,
41}
42
43impl<F: PrimeField, H: Hash<F>> Node<F, H> {
44 pub fn new_leaf(v: F) -> Self {
45 let h = H::hash(&[v]);
46 Self {
47 phantom: PhantomData::<H>,
48 hash: h,
49 left: None,
50 right: None,
51 value: Some(v),
52 }
53 }
54 pub fn new_node(l: Self, r: Self) -> Self {
55 let left = Box::new(l);
56 let right = Box::new(r);
57 let hash = H::hash(&[left.hash, right.hash]);
58 Self {
59 phantom: PhantomData::<H>,
60 hash,
61 left: Some(left),
62 right: Some(right),
63 value: None,
64 }
65 }
66}
67
68pub struct MerkleTree<F: PrimeField, H: Hash<F>> {
69 pub root: Node<F, H>,
70 nlevels: u32,
71}
72
73impl<F: PrimeField, H: Hash<F>> MerkleTree<F, H> {
74 pub fn commit(values: &[F]) -> (F, Self) {
75 // for the moment assume that values length is a power of 2.
76 if (values.len() != 0) && (values.len() & (values.len() - 1) != 0) {
77 panic!("values.len() should be a power of 2");
78 }
79
80 // prepare the leafs
81 let mut leaf_nodes: Vec<Node<F, H>> = Vec::new();
82 for i in 0..values.len() {
83 let node = Node::<F, H>::new_leaf(values[i]);
84 leaf_nodes.push(node);
85 }
86 // go up from the leafs to the root
87 let top_nodes = Self::up_from_nodes(leaf_nodes);
88
89 (
90 top_nodes[0].hash,
91 Self {
92 root: top_nodes[0].clone(),
93 nlevels: log2(values.len()),
94 },
95 )
96 }
97 fn up_from_nodes(nodes: Vec<Node<F, H>>) -> Vec<Node<F, H>> {
98 if nodes.len() == 0 {
99 return [Node::<F, H> {
100 phantom: PhantomData::<H>,
101 hash: F::from(0_u32),
102 left: None,
103 right: None,
104 value: None,
105 }]
106 .to_vec();
107 }
108 if nodes.len() == 1 {
109 return nodes;
110 }
111 let mut next_level_nodes: Vec<Node<F, H>> = Vec::new();
112 for i in (0..nodes.len()).step_by(2) {
113 let node = Node::<F, H>::new_node(nodes[i].clone(), nodes[i + 1].clone());
114 next_level_nodes.push(node);
115 }
116 return Self::up_from_nodes(next_level_nodes);
117 }
118 fn get_path(num_levels: u32, value: F) -> Vec<bool> {
119 let value_bytes = value.into_bigint().to_bytes_le();
120 let mut path = Vec::new();
121 for i in 0..num_levels {
122 path.push(value_bytes[(i / 8) as usize] & (1 << (i % 8)) != 0);
123 }
124 path
125 }
126 pub fn open(&self, index: F) -> Vec<F> {
127 // start from root, and go down to the index, while getting the siblings at each level
128 let path = Self::get_path(self.nlevels, index);
129 // reverse path as we're going from up to down
130 let path_inv = path.iter().copied().rev().collect();
131 let mut siblings: Vec<F> = Vec::new();
132 siblings = Self::go_down(path_inv, self.root.clone(), siblings);
133 return siblings;
134 }
135 fn go_down(path: Vec<bool>, node: Node<F, H>, mut siblings: Vec<F>) -> Vec<F> {
136 if !node.value.is_none() {
137 return siblings;
138 }
139 if !path[0] {
140 siblings.push(node.right.unwrap().hash);
141 return Self::go_down(path[1..].to_vec(), *node.left.unwrap(), siblings);
142 } else {
143 siblings.push(node.left.unwrap().hash);
144 return Self::go_down(path[1..].to_vec(), *node.right.unwrap(), siblings);
145 }
146 }
147 pub fn verify(root: F, index: F, value: F, siblings: Vec<F>) -> bool {
148 let mut h = H::hash(&[value]);
149 let path = Self::get_path(siblings.len() as u32, index);
150 for i in 0..siblings.len() {
151 if !path[i] {
152 h = H::hash(&[h, siblings[siblings.len() - 1 - i]]);
153 } else {
154 h = H::hash(&[siblings[siblings.len() - 1 - i], h]);
155 }
156 }
157 if h == root {
158 return true;
159 }
160 false
161 }
162}
163
164#[cfg(test)]
165mod tests {
166 use super::*;
167 use ark_std::UniformRand;
168 pub type Fr = ark_bn254::Fr; // scalar field
169
170 #[test]
171 fn test_path() {
172 assert_eq!(
173 MerkleTree::<Fr, Keccak256Hash<Fr>>::get_path(8, Fr::from(0_u32)),
174 [false, false, false, false, false, false, false, false]
175 );
176 assert_eq!(
177 MerkleTree::<Fr, Keccak256Hash<Fr>>::get_path(8, Fr::from(1_u32)),
178 [true, false, false, false, false, false, false, false]
179 );
180 assert_eq!(
181 MerkleTree::<Fr, Keccak256Hash<Fr>>::get_path(8, Fr::from(2_u32)),
182 [false, true, false, false, false, false, false, false]
183 );
184 assert_eq!(
185 MerkleTree::<Fr, Keccak256Hash<Fr>>::get_path(8, Fr::from(3_u32)),
186 [true, true, false, false, false, false, false, false]
187 );
188 assert_eq!(
189 MerkleTree::<Fr, Keccak256Hash<Fr>>::get_path(8, Fr::from(254_u32)),
190 [false, true, true, true, true, true, true, true]
191 );
192 assert_eq!(
193 MerkleTree::<Fr, Keccak256Hash<Fr>>::get_path(8, Fr::from(255_u32)),
194 [true, true, true, true, true, true, true, true]
195 );
196 }
197
198 #[test]
199 fn test_new_empty_tree() {
200 let (root, mt) = MerkleTree::<Fr, Keccak256Hash<Fr>>::commit(&[]);
201 assert_eq!(mt.root.hash, Fr::from(0_u32));
202 assert_eq!(root, Fr::from(0_u32));
203 }
204
205 #[test]
206 fn test_proof() {
207 type MT = MerkleTree<Fr, Keccak256Hash<Fr>>;
208
209 let values = [
210 Fr::from(0_u32),
211 Fr::from(1_u32),
212 Fr::from(2_u32),
213 Fr::from(3_u32),
214 Fr::from(200_u32),
215 Fr::from(201_u32),
216 Fr::from(202_u32),
217 Fr::from(203_u32),
218 ];
219
220 let (root, mt) = MT::commit(&values);
221 assert_eq!(
222 root.to_string(),
223 "6195952497672867974990959901930625199810318409246598214215524466855665265259"
224 );
225
226 let index = 3;
227 let index_F = Fr::from(index as u32);
228 let siblings = mt.open(index_F);
229
230 assert!(MT::verify(root, index_F, values[index], siblings));
231 }
232
233 #[test]
234 fn test_proofs() {
235 type MT = MerkleTree<Fr, Keccak256Hash<Fr>>;
236 let mut rng = ark_std::test_rng();
237
238 let n_values = 64;
239 let mut values: Vec<Fr> = Vec::new();
240 for _i in 0..n_values {
241 let v = Fr::rand(&mut rng);
242 values.push(v);
243 }
244
245 let (root, mt) = MT::commit(&values);
246 assert_eq!(mt.nlevels, 6);
247
248 for i in 0..n_values {
249 let i_Fr = Fr::from(i as u32);
250 let siblings = mt.open(i_Fr);
251 assert!(MT::verify(root, i_Fr, values[i], siblings));
252 }
253 }
254}