1use std::{env, fs::File, time::Instant};
  2
  3use ark_bls12_381::{Bls12_381, Fr};
  4use ark_serialize::Write;
  5use ark_std::test_rng;
  6use hyperplonk::{
  7    prelude::{CustomizedGates, HyperPlonkErrors, MockCircuit},
  8    HyperPlonkSNARK,
  9};
 10use rayon::ThreadPoolBuilder;
 11use subroutines::{
 12    pcs::{
 13        prelude::{MultilinearKzgPCS, MultilinearUniversalParams},
 14        PolynomialCommitmentScheme,
 15    },
 16    poly_iop::PolyIOP,
 17};
 18
 19const SUPPORTED_SIZE: usize = 15;
 20const MIN_NUM_VARS: usize = 8;
 21const MAX_NUM_VARS: usize = 15;
 22const MIN_CUSTOM_DEGREE: usize = 1;
 23const MAX_CUSTOM_DEGREE: usize = 32;
 24const MAX_NUM_THREADS: usize = 24;
 25
 26fn main() -> Result<(), HyperPlonkErrors> {
 27    let args: Vec<String> = env::args().collect();
 28    let thread = args[1].parse().unwrap_or(MAX_NUM_THREADS);
 29    let mut rng = test_rng();
 30    let pcs_srs = MultilinearKzgPCS::<Bls12_381>::gen_srs_for_testing(&mut rng, SUPPORTED_SIZE)?;
 31    ThreadPoolBuilder::new()
 32        .num_threads(thread)
 33        .build()
 34        .unwrap();
 35    bench_vanilla_plonk(&pcs_srs, thread)?;
 36    for degree in MIN_CUSTOM_DEGREE..MAX_CUSTOM_DEGREE {
 37        bench_high_degree_plonk(&pcs_srs, degree, thread)?;
 38    }
 39
 40    Ok(())
 41}
 42
 43fn bench_vanilla_plonk(
 44    pcs_srs: &MultilinearUniversalParams<Bls12_381>,
 45    thread: usize,
 46) -> Result<(), HyperPlonkErrors> {
 47    let filename = format!("vanilla threads {}.txt", thread);
 48    let mut file = File::create(filename).unwrap();
 49    for nv in MIN_NUM_VARS..MAX_NUM_VARS {
 50        let vanilla_gate = CustomizedGates::vanilla_plonk_gate();
 51        bench_mock_circuit_zkp_helper(&mut file, nv, &vanilla_gate, &pcs_srs)?;
 52    }
 53
 54    Ok(())
 55}
 56
 57fn bench_high_degree_plonk(
 58    pcs_srs: &MultilinearUniversalParams<Bls12_381>,
 59    degree: usize,
 60    thread: usize,
 61) -> Result<(), HyperPlonkErrors> {
 62    let filename = format!("high degree {} thread {}.txt", degree, thread);
 63    let mut file = File::create(filename).unwrap();
 64    for nv in MIN_NUM_VARS..MAX_NUM_VARS {
 65        let vanilla_gate = CustomizedGates::mock_gate(2, degree);
 66        bench_mock_circuit_zkp_helper(&mut file, nv, &vanilla_gate, &pcs_srs)?;
 67    }
 68
 69    Ok(())
 70}
 71
 72fn bench_mock_circuit_zkp_helper(
 73    file: &mut File,
 74    nv: usize,
 75    gate: &CustomizedGates,
 76    pcs_srs: &MultilinearUniversalParams<Bls12_381>,
 77) -> Result<(), HyperPlonkErrors> {
 78    let repetition = if nv < 10 {
 79        5
 80    } else if nv < 20 {
 81        2
 82    } else {
 83        1
 84    };
 85
 86    //==========================================================
 87    let start = Instant::now();
 88    for _ in 0..repetition {
 89        let circuit = MockCircuit::<Fr>::new(1 << nv, gate);
 90        assert!(circuit.is_satisfied());
 91    }
 92    println!(
 93        "mock circuit gen for {} variables: {} ns",
 94        nv,
 95        start.elapsed().as_nanos() / repetition as u128
 96    );
 97
 98    let circuit = MockCircuit::<Fr>::new(1 << nv, gate);
 99    assert!(circuit.is_satisfied());
100    let index = circuit.index;
101    //==========================================================
102    // generate pk and vks
103    let start = Instant::now();
104    for _ in 0..repetition {
105        let (_pk, _vk) = <PolyIOP<Fr> as HyperPlonkSNARK<
106            Bls12_381,
107            MultilinearKzgPCS<Bls12_381>,
108        >>::preprocess(&index, &pcs_srs)?;
109    }
110    println!(
111        "key extraction for {} variables: {} us",
112        nv,
113        start.elapsed().as_micros() / repetition as u128
114    );
115    let (pk, vk) =
116        <PolyIOP<Fr> as HyperPlonkSNARK<Bls12_381, MultilinearKzgPCS<Bls12_381>>>::preprocess(
117            &index, &pcs_srs,
118        )?;
119    //==========================================================
120    // generate a proof
121    let start = Instant::now();
122    for _ in 0..repetition {
123        let _proof =
124            <PolyIOP<Fr> as HyperPlonkSNARK<Bls12_381, MultilinearKzgPCS<Bls12_381>>>::prove(
125                &pk,
126                &circuit.witnesses[0].coeff_ref(),
127                &circuit.witnesses,
128            )?;
129    }
130    let t = start.elapsed().as_micros() / repetition as u128;
131
132    file.write_all(format!("{} {}\n", nv, t).as_ref()).unwrap();
133
134    let proof = <PolyIOP<Fr> as HyperPlonkSNARK<Bls12_381, MultilinearKzgPCS<Bls12_381>>>::prove(
135        &pk,
136        &circuit.witnesses[0].coeff_ref(),
137        &circuit.witnesses,
138    )?;
139    //==========================================================
140    // verify a proof
141    let start = Instant::now();
142    for _ in 0..repetition {
143        let verify =
144            <PolyIOP<Fr> as HyperPlonkSNARK<Bls12_381, MultilinearKzgPCS<Bls12_381>>>::verify(
145                &vk,
146                &circuit.witnesses[0].coeff_ref(),
147                &proof,
148            )?;
149        assert!(verify);
150    }
151    println!(
152        "verifying for {} variables: {} us",
153        nv,
154        start.elapsed().as_micros() / repetition as u128
155    );
156    Ok(())
157}