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 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 for leaf in &padded_leaves {
31 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 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 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 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