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_std::log2;
  8use arkworks_native_gadgets::poseidon;
  9use arkworks_native_gadgets::poseidon::FieldHasher;
 10use arkworks_utils::{
 11    bytes_matrix_to_f, bytes_vec_to_f, poseidon_params::setup_poseidon_params, Curve,
 12};
 13
 14pub struct Params<F: PrimeField> {
 15    pub poseidon_hash: poseidon::Poseidon<F>,
 16}
 17
 18#[derive(Clone, Debug)]
 19pub struct Node<F: PrimeField> {
 20    hash: F,
 21    left: Option<Box<Node<F>>>,
 22    right: Option<Box<Node<F>>>,
 23    value: Option<F>,
 24}
 25
 26impl<F: PrimeField> Node<F> {
 27    pub fn new_leaf(params: &Params<F>, v: F) -> Self {
 28        let h = params.poseidon_hash.hash(&[v]).unwrap();
 29        Self {
 30            hash: h,
 31            left: None,
 32            right: None,
 33            value: Some(v),
 34        }
 35    }
 36    pub fn new_node(params: &Params<F>, l: Self, r: Self) -> Self {
 37        let left = Box::new(l);
 38        let right = Box::new(r);
 39        let hash = params.poseidon_hash.hash(&[left.hash, right.hash]).unwrap();
 40        Self {
 41            hash,
 42            left: Some(left),
 43            right: Some(right),
 44            value: None,
 45        }
 46    }
 47}
 48
 49pub struct MerkleTree<F: PrimeField> {
 50    pub root: Node<F>,
 51    nlevels: u32,
 52}
 53
 54impl<F: PrimeField> MerkleTree<F> {
 55    pub fn setup(poseidon_hash: &poseidon::Poseidon<F>) -> Params<F> {
 56        Params {
 57            poseidon_hash: poseidon_hash.clone(),
 58        }
 59    }
 60    pub fn new(params: &Params<F>, values: Vec<F>) -> Self {
 61        // for the moment assume that values length is a power of 2.
 62        if (values.len() != 0) && (values.len() & (values.len() - 1) != 0) {
 63            panic!("values.len() should be a power of 2");
 64        }
 65
 66        // prepare the leafs
 67        let mut leaf_nodes: Vec<Node<F>> = Vec::new();
 68        for i in 0..values.len() {
 69            let node = Node::<F>::new_leaf(&params, values[i]);
 70            leaf_nodes.push(node);
 71        }
 72        // go up from the leafs to the root
 73        let top_nodes = Self::up_from_nodes(&params, leaf_nodes);
 74
 75        Self {
 76            root: top_nodes[0].clone(),
 77            nlevels: log2(values.len()),
 78        }
 79    }
 80    fn up_from_nodes(params: &Params<F>, nodes: Vec<Node<F>>) -> Vec<Node<F>> {
 81        if nodes.len() == 0 {
 82            return [Node::<F> {
 83                hash: F::from(0_u32),
 84                left: None,
 85                right: None,
 86                value: None,
 87            }]
 88            .to_vec();
 89        }
 90        if nodes.len() == 1 {
 91            return nodes;
 92        }
 93        let mut next_level_nodes: Vec<Node<F>> = Vec::new();
 94        for i in (0..nodes.len()).step_by(2) {
 95            let node = Node::<F>::new_node(&params, nodes[i].clone(), nodes[i + 1].clone());
 96            next_level_nodes.push(node);
 97        }
 98        return Self::up_from_nodes(params, next_level_nodes);
 99    }
100    fn get_path(num_levels: u32, value: F) -> Vec<bool> {
101        let value_bytes = value.into_repr().to_bytes_le();
102        let mut path = Vec::new();
103        for i in 0..num_levels {
104            path.push(value_bytes[(i / 8) as usize] & (1 << (i % 8)) != 0);
105        }
106        path
107    }
108    pub fn gen_proof(&self, index: F) -> Vec<F> {
109        // start from root, and go down to the index, while getting the siblings at each level
110        let path = Self::get_path(self.nlevels, index);
111        // reverse path as we're going from up to down
112        let path_inv = path.iter().copied().rev().collect();
113        let mut siblings: Vec<F> = Vec::new();
114        siblings = Self::go_down(path_inv, self.root.clone(), siblings);
115        return siblings;
116    }
117    fn go_down(path: Vec<bool>, node: Node<F>, mut siblings: Vec<F>) -> Vec<F> {
118        if !node.value.is_none() {
119            return siblings;
120        }
121        if !path[0] {
122            siblings.push(node.right.unwrap().hash);
123            return Self::go_down(path[1..].to_vec(), *node.left.unwrap(), siblings);
124        } else {
125            siblings.push(node.left.unwrap().hash);
126            return Self::go_down(path[1..].to_vec(), *node.right.unwrap(), siblings);
127        }
128    }
129    pub fn verify(params: &Params<F>, root: F, index: F, value: F, siblings: Vec<F>) -> bool {
130        let mut h = params.poseidon_hash.hash(&[value]).unwrap();
131        let path = Self::get_path(siblings.len() as u32, index);
132        for i in 0..siblings.len() {
133            if !path[i] {
134                h = params
135                    .poseidon_hash
136                    .hash(&[h, siblings[siblings.len() - 1 - i]])
137                    .unwrap();
138            } else {
139                h = params
140                    .poseidon_hash
141                    .hash(&[siblings[siblings.len() - 1 - i], h])
142                    .unwrap();
143            }
144        }
145        if h == root {
146            return true;
147        }
148        false
149    }
150}
151
152pub struct MerkleTreePoseidon<F: PrimeField>(MerkleTree<F>);
153
154pub struct MTProof<F: PrimeField> {
155    index: F,
156    siblings: Vec<F>,
157}
158
159impl<F: PrimeField> MerkleTreePoseidon<F> {
160    pub fn commit(values: &[F]) -> (F, Self) {
161        let poseidon_params = poseidon_setup_params::<F>(Curve::Bn254, 5, 4);
162        let poseidon_hash = poseidon::Poseidon::new(poseidon_params);
163        let params = MerkleTree::setup(&poseidon_hash);
164        let mt = MerkleTree::new(&params, values.to_vec());
165        (mt.root.hash, MerkleTreePoseidon(mt))
166    }
167    pub fn open(&self, index: F) -> Vec<F> {
168        self.0.gen_proof(index)
169    }
170    pub fn verify(root: F, index: F, value: F, siblings: Vec<F>) -> bool {
171        let poseidon_params = poseidon_setup_params::<F>(Curve::Bn254, 5, 4);
172        let poseidon_hash = poseidon::Poseidon::new(poseidon_params);
173        let params = MerkleTree::setup(&poseidon_hash);
174        MerkleTree::verify(&params, root, index, value, siblings)
175    }
176}
177
178pub fn poseidon_setup_params<F: PrimeField>(
179    curve: Curve,
180    exp: i8,
181    width: u8,
182) -> poseidon::PoseidonParameters<F> {
183    let pos_data = setup_poseidon_params(curve, exp, width).unwrap();
184
185    let mds_f = bytes_matrix_to_f(&pos_data.mds);
186    let rounds_f = bytes_vec_to_f(&pos_data.rounds);
187
188    poseidon::PoseidonParameters {
189        mds_matrix: mds_f,
190        round_keys: rounds_f,
191        full_rounds: pos_data.full_rounds,
192        partial_rounds: pos_data.partial_rounds,
193        sbox: poseidon::sbox::PoseidonSbox(pos_data.exp),
194        width: pos_data.width,
195    }
196}
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201    use ark_std::UniformRand;
202    pub type Fr = ark_bn254::Fr; // scalar field
203
204    #[test]
205    fn test_path() {
206        assert_eq!(
207            MerkleTree::get_path(8, Fr::from(0_u32)),
208            [false, false, false, false, false, false, false, false]
209        );
210        assert_eq!(
211            MerkleTree::get_path(8, Fr::from(1_u32)),
212            [true, false, false, false, false, false, false, false]
213        );
214        assert_eq!(
215            MerkleTree::get_path(8, Fr::from(2_u32)),
216            [false, true, false, false, false, false, false, false]
217        );
218        assert_eq!(
219            MerkleTree::get_path(8, Fr::from(3_u32)),
220            [true, true, false, false, false, false, false, false]
221        );
222        assert_eq!(
223            MerkleTree::get_path(8, Fr::from(254_u32)),
224            [false, true, true, true, true, true, true, true]
225        );
226        assert_eq!(
227            MerkleTree::get_path(8, Fr::from(255_u32)),
228            [true, true, true, true, true, true, true, true]
229        );
230    }
231
232    #[test]
233    fn test_new_empty_tree() {
234        let poseidon_params = poseidon_setup_params::<Fr>(Curve::Bn254, 5, 4);
235        let poseidon_hash = poseidon::Poseidon::new(poseidon_params);
236
237        let params = MerkleTree::setup(&poseidon_hash);
238        let mt = MerkleTree::new(&params, Vec::new());
239        assert_eq!(mt.root.hash, Fr::from(0_u32));
240    }
241
242    #[test]
243    fn test_proof() {
244        let poseidon_params = poseidon_setup_params::<Fr>(Curve::Bn254, 5, 4);
245        let poseidon_hash = poseidon::Poseidon::new(poseidon_params);
246
247        let params = MerkleTree::setup(&poseidon_hash);
248
249        let values = [
250            Fr::from(0_u32),
251            Fr::from(1_u32),
252            Fr::from(2_u32),
253            Fr::from(3_u32),
254            Fr::from(200_u32),
255            Fr::from(201_u32),
256            Fr::from(202_u32),
257            Fr::from(203_u32),
258        ];
259
260        let mt = MerkleTree::new(&params, values.to_vec());
261        assert_eq!(
262            mt.root.hash.to_string(),
263            "Fp256 \"(10E90845D7A113686E4F2F30D73B315BA891A5DB9A58782F1260C35F99660794)\""
264        );
265
266        let index = 3;
267        let index_F = Fr::from(index as u32);
268        let siblings = mt.gen_proof(index_F);
269
270        assert!(MerkleTree::verify(
271            &params,
272            mt.root.hash,
273            index_F,
274            values[index],
275            siblings
276        ));
277    }
278
279    #[test]
280    fn test_proofs() {
281        let poseidon_params = poseidon_setup_params::<Fr>(Curve::Bn254, 5, 4);
282        let poseidon_hash = poseidon::Poseidon::new(poseidon_params);
283
284        let params = MerkleTree::setup(&poseidon_hash);
285
286        let mut rng = ark_std::test_rng();
287
288        let n_values = 64;
289        let mut values: Vec<Fr> = Vec::new();
290        for _i in 0..n_values {
291            let v = Fr::rand(&mut rng);
292            values.push(v);
293        }
294
295        let mt = MerkleTree::new(&params, values.to_vec());
296        assert_eq!(mt.nlevels, 6);
297
298        for i in 0..n_values {
299            let i_Fr = Fr::from(i as u32);
300            let siblings = mt.gen_proof(i_Fr);
301            assert!(MerkleTree::verify(
302                &params,
303                mt.root.hash,
304                i_Fr,
305                values[i],
306                siblings
307            ));
308        }
309    }
310}