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: usize) -> 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, F::from(index as u32));
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: usize, 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, F::from(index as u32));
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>);
153impl<F: PrimeField> MerkleTreePoseidon<F> {
154    pub fn commit(values: &[F]) -> (F, Self) {
155        let poseidon_params = poseidon_setup_params::<F>(Curve::Bn254, 5, 4);
156        let poseidon_hash = poseidon::Poseidon::new(poseidon_params);
157        let params = MerkleTree::setup(&poseidon_hash);
158        let mt = MerkleTree::new(&params, values.to_vec());
159        (mt.root.hash, MerkleTreePoseidon(mt))
160    }
161    pub fn prove(&self, index: usize) -> Vec<F> {
162        self.0.gen_proof(index)
163    }
164    pub fn verify(root: F, index: usize, value: F, siblings: Vec<F>) -> bool {
165        let poseidon_params = poseidon_setup_params::<F>(Curve::Bn254, 5, 4);
166        let poseidon_hash = poseidon::Poseidon::new(poseidon_params);
167        let params = MerkleTree::setup(&poseidon_hash);
168        MerkleTree::verify(&params, root, index, value, siblings)
169    }
170}
171
172pub fn poseidon_setup_params<F: PrimeField>(
173    curve: Curve,
174    exp: i8,
175    width: u8,
176) -> poseidon::PoseidonParameters<F> {
177    let pos_data = setup_poseidon_params(curve, exp, width).unwrap();
178
179    let mds_f = bytes_matrix_to_f(&pos_data.mds);
180    let rounds_f = bytes_vec_to_f(&pos_data.rounds);
181
182    poseidon::PoseidonParameters {
183        mds_matrix: mds_f,
184        round_keys: rounds_f,
185        full_rounds: pos_data.full_rounds,
186        partial_rounds: pos_data.partial_rounds,
187        sbox: poseidon::sbox::PoseidonSbox(pos_data.exp),
188        width: pos_data.width,
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195    use ark_std::UniformRand;
196    pub type Fr = ark_bn254::Fr; // scalar field
197
198    #[test]
199    fn test_path() {
200        assert_eq!(
201            MerkleTree::get_path(8, Fr::from(0_u32)),
202            [false, false, false, false, false, false, false, false]
203        );
204        assert_eq!(
205            MerkleTree::get_path(8, Fr::from(1_u32)),
206            [true, false, false, false, false, false, false, false]
207        );
208        assert_eq!(
209            MerkleTree::get_path(8, Fr::from(2_u32)),
210            [false, true, false, false, false, false, false, false]
211        );
212        assert_eq!(
213            MerkleTree::get_path(8, Fr::from(3_u32)),
214            [true, true, false, false, false, false, false, false]
215        );
216        assert_eq!(
217            MerkleTree::get_path(8, Fr::from(254_u32)),
218            [false, true, true, true, true, true, true, true]
219        );
220        assert_eq!(
221            MerkleTree::get_path(8, Fr::from(255_u32)),
222            [true, true, true, true, true, true, true, true]
223        );
224    }
225
226    #[test]
227    fn test_new_empty_tree() {
228        let poseidon_params = poseidon_setup_params::<Fr>(Curve::Bn254, 5, 4);
229        let poseidon_hash = poseidon::Poseidon::new(poseidon_params);
230
231        let params = MerkleTree::setup(&poseidon_hash);
232        let mt = MerkleTree::new(&params, Vec::new());
233        assert_eq!(mt.root.hash, Fr::from(0_u32));
234    }
235
236    #[test]
237    fn test_proof() {
238        let poseidon_params = poseidon_setup_params::<Fr>(Curve::Bn254, 5, 4);
239        let poseidon_hash = poseidon::Poseidon::new(poseidon_params);
240
241        let params = MerkleTree::setup(&poseidon_hash);
242
243        let values = [
244            Fr::from(0_u32),
245            Fr::from(1_u32),
246            Fr::from(2_u32),
247            Fr::from(3_u32),
248            Fr::from(200_u32),
249            Fr::from(201_u32),
250            Fr::from(202_u32),
251            Fr::from(203_u32),
252        ];
253
254        let mt = MerkleTree::new(&params, values.to_vec());
255        assert_eq!(
256            mt.root.hash.to_string(),
257            "Fp256 \"(10E90845D7A113686E4F2F30D73B315BA891A5DB9A58782F1260C35F99660794)\""
258        );
259
260        let index = 3;
261        let siblings = mt.gen_proof(index);
262
263        assert!(MerkleTree::verify(
264            &params,
265            mt.root.hash,
266            index,
267            values[index],
268            siblings
269        ));
270    }
271
272    #[test]
273    fn test_proofs() {
274        let poseidon_params = poseidon_setup_params::<Fr>(Curve::Bn254, 5, 4);
275        let poseidon_hash = poseidon::Poseidon::new(poseidon_params);
276
277        let params = MerkleTree::setup(&poseidon_hash);
278
279        let mut rng = ark_std::test_rng();
280
281        let n_values = 256;
282        let mut values: Vec<Fr> = Vec::new();
283        for _i in 0..n_values {
284            let v = Fr::rand(&mut rng);
285            values.push(v);
286        }
287
288        let mt = MerkleTree::new(&params, values.to_vec());
289        assert_eq!(mt.nlevels, 8);
290
291        for i in 0..n_values {
292            let siblings = mt.gen_proof(i);
293            assert!(MerkleTree::verify(
294                &params,
295                mt.root.hash,
296                i,
297                values[i],
298                siblings
299            ));
300        }
301    }
302}