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