merkle_tree_wasm/
lib.rs

1#![allow(non_snake_case)]
2mod merkle_tree_wasm;
3
4extern crate console_error_panic_hook;
5
6use ark_ff::{BigInteger, PrimeField};
7use ark_serialize::CanonicalSerialize;
8use num_bigint::BigUint;
9use anonklub_poseidon::constants::secp256k1_w3;
10use wasm_bindgen::{prelude::wasm_bindgen, JsValue};
11
12pub use merkle_tree_wasm::{MerkleProofBytes, MerkleTree};
13
14fn internal_generate_merkle_proof<F: PrimeField>(
15    leaves: Vec<String>,
16    leaf: String,
17    depth: usize,
18) -> Result<MerkleProofBytes, String> {
19    let mut padded_leaves = leaves.clone();
20    // Pad the leaves to equal the size of the tree
21    // Needs to be an even string
22    padded_leaves.resize(1 << depth, "00".to_string());
23
24    const ARITY: usize = 2;
25    const WIDTH: usize = ARITY + 1;
26
27    let mut tree = MerkleTree::<F, WIDTH>::new(secp256k1_w3());
28
29    // Insert all the leaves into the tree
30    for leaf in &padded_leaves {
31        // Converting String to F
32        let leaf_hex = hex::decode(leaf.replace("0x", ""))
33            .map_err(|e| format!("MerkleTree: Error decoding hex: {}", e))?;
34        let leaf_bytes = F::from(BigUint::from_bytes_be(&leaf_hex));
35        tree.insert(leaf_bytes);
36    }
37
38    tree.finish();
39
40    let leaf_hex = hex::decode(leaf.replace("0x", ""))
41        .map_err(|e| format!("MerkleTree: Error decoding hex: {}", e))?;
42    let leaf_prime_field = F::from(BigUint::from_bytes_be(&leaf_hex));
43
44    let proof = tree.create_proof(leaf_prime_field)?;
45
46    // Encode the Merkle Proof output with BigEndian
47    let mut merkle_siblings = Vec::with_capacity(depth);
48    let mut merkle_indices = Vec::with_capacity(depth);
49    let siblings_bytes = proof
50        .siblings
51        .iter()
52        .flat_map(|sibling| sibling.into_bigint().to_bytes_be())
53        .collect::<Vec<u8>>();
54
55    let indices_bytes = proof
56        .path_indices
57        .iter()
58        .flat_map(|i| F::from(*i as u32).into_bigint().to_bytes_be())
59        .collect::<Vec<u8>>();
60
61    merkle_siblings.extend_from_slice(&siblings_bytes);
62    merkle_indices.extend_from_slice(&indices_bytes);
63
64    let root_bytes = tree
65        .root
66        .ok_or("MerkleTree: Root is not available")?
67        .into_bigint()
68        .to_bytes_be();
69
70    Ok(MerkleProofBytes {
71        siblings: siblings_bytes,
72        path_indices: indices_bytes,
73        root: root_bytes,
74    })
75}
76
77#[cfg(target_arch = "wasm32")]
78#[wasm_bindgen]
79pub fn generate_merkle_proof(
80    leaves: Vec<String>,
81    leaf: String,
82    depth: usize,
83) -> Result<Vec<u8>, JsValue> {
84    std::panic::set_hook(Box::new(console_error_panic_hook::hook));
85
86    type F = ark_secp256k1::Fq;
87
88    let merkle_proof_bytes = internal_generate_merkle_proof::<F>(leaves, leaf, depth)
89        .map_err(|e| JsValue::from_str(&e))?;
90
91    // Serialize the full merkle proof
92    let mut merkle_proof_bytes_serialized = Vec::new();
93    merkle_proof_bytes
94        .serialize_compressed(&mut merkle_proof_bytes_serialized)
95        .map_err(|_e| JsValue::from_str("Error serializing Merkle proof bytes"))?;
96
97    Ok(merkle_proof_bytes_serialized)
98}
99
100#[cfg(not(target_arch = "wasm32"))]
101pub fn generate_merkle_proof(
102    leaves: Vec<String>,
103    leaf: String,
104    depth: usize,
105) -> Result<Vec<u8>, String> {
106    type F = ark_secp256k1::Fq;
107
108    let merkle_proof_bytes = internal_generate_merkle_proof::<F>(leaves, leaf, depth)?;
109
110    // Serialize the full merkle proof
111    let mut merkle_proof_bytes_serialized = Vec::new();
112    merkle_proof_bytes
113        .serialize_compressed(&mut merkle_proof_bytes_serialized)
114        .map_err(|e| format!("Error serializing Merkle proof bytes: {}", e))?;
115
116    Ok(merkle_proof_bytes_serialized)
117}
118
119// TODO: writing tests