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