1use std::{fs::File, io, time::Instant};
  2
  3use ark_bls12_381::{Bls12_381, Fr};
  4use ark_serialize::{CanonicalDeserialize, CanonicalSerialize, Write};
  5use ark_std::test_rng;
  6use hyperplonk::{
  7    prelude::{CustomizedGates, HyperPlonkErrors, MockCircuit},
  8    HyperPlonkSNARK,
  9};
 10use subroutines::{
 11    pcs::{
 12        prelude::{MultilinearKzgPCS, MultilinearUniversalParams},
 13        PolynomialCommitmentScheme,
 14    },
 15    poly_iop::PolyIOP,
 16};
 17
 18const SUPPORTED_SIZE: usize = 20;
 19const MIN_NUM_VARS: usize = 8;
 20const MAX_NUM_VARS: usize = 20;
 21const MIN_CUSTOM_DEGREE: usize = 1;
 22const MAX_CUSTOM_DEGREE: usize = 32;
 23const HIGH_DEGREE_TEST_NV: usize = 15;
 24
 25fn main() -> Result<(), HyperPlonkErrors> {
 26    let thread = rayon::current_num_threads();
 27    println!("start benchmark with #{} threads", thread);
 28    let mut rng = test_rng();
 29    let pcs_srs = {
 30        match read_srs() {
 31            Ok(p) => p,
 32            Err(_e) => {
 33                let srs =
 34                    MultilinearKzgPCS::<Bls12_381>::gen_srs_for_testing(&mut rng, SUPPORTED_SIZE)?;
 35                write_srs(&srs);
 36                srs
 37            },
 38        }
 39    };
 40    // bench_jellyfish_plonk(&pcs_srs, thread)?;
 41    // println!();
 42    bench_vanilla_plonk(&pcs_srs, thread)?;
 43    println!();
 44    // for degree in MIN_CUSTOM_DEGREE..=MAX_CUSTOM_DEGREE {
 45    //     bench_high_degree_plonk(&pcs_srs, degree, thread)?;
 46    //     println!();
 47    // }
 48    // println!();
 49
 50    Ok(())
 51}
 52
 53fn read_srs() -> Result<MultilinearUniversalParams<Bls12_381>, io::Error> {
 54    let mut f = File::open("srs.params")?;
 55    Ok(MultilinearUniversalParams::<Bls12_381>::deserialize_unchecked(&mut f).unwrap())
 56}
 57
 58fn write_srs(pcs_srs: &MultilinearUniversalParams<Bls12_381>) {
 59    let mut f = File::create("srs.params").unwrap();
 60    pcs_srs.serialize_uncompressed(&mut f).unwrap();
 61}
 62
 63fn bench_vanilla_plonk(
 64    pcs_srs: &MultilinearUniversalParams<Bls12_381>,
 65    thread: usize,
 66) -> Result<(), HyperPlonkErrors> {
 67    let filename = format!("vanilla threads {}.txt", thread);
 68    let mut file = File::create(filename).unwrap();
 69    for nv in MIN_NUM_VARS..=MAX_NUM_VARS {
 70        let vanilla_gate = CustomizedGates::vanilla_plonk_gate();
 71        bench_mock_circuit_zkp_helper(&mut file, nv, &vanilla_gate, &pcs_srs)?;
 72    }
 73
 74    Ok(())
 75}
 76
 77fn bench_jellyfish_plonk(
 78    pcs_srs: &MultilinearUniversalParams<Bls12_381>,
 79    thread: usize,
 80) -> Result<(), HyperPlonkErrors> {
 81    let filename = format!("jellyfish threads {}.txt", thread);
 82    let mut file = File::create(filename).unwrap();
 83    for nv in MIN_NUM_VARS..=MAX_NUM_VARS {
 84        let jf_gate = CustomizedGates::jellyfish_turbo_plonk_gate();
 85        bench_mock_circuit_zkp_helper(&mut file, nv, &jf_gate, &pcs_srs)?;
 86    }
 87
 88    Ok(())
 89}
 90
 91fn bench_high_degree_plonk(
 92    pcs_srs: &MultilinearUniversalParams<Bls12_381>,
 93    degree: usize,
 94    thread: usize,
 95) -> Result<(), HyperPlonkErrors> {
 96    let filename = format!("high degree {} thread {}.txt", degree, thread);
 97    let mut file = File::create(filename).unwrap();
 98    println!("custom gate of degree {}", degree);
 99    let vanilla_gate = CustomizedGates::mock_gate(2, degree);
100    bench_mock_circuit_zkp_helper(&mut file, HIGH_DEGREE_TEST_NV, &vanilla_gate, &pcs_srs)?;
101
102    Ok(())
103}
104
105fn bench_mock_circuit_zkp_helper(
106    file: &mut File,
107    nv: usize,
108    gate: &CustomizedGates,
109    pcs_srs: &MultilinearUniversalParams<Bls12_381>,
110) -> Result<(), HyperPlonkErrors> {
111    let repetition = if nv < 10 {
112        5
113    } else if nv < 20 {
114        2
115    } else {
116        1
117    };
118
119    //==========================================================
120    let circuit = MockCircuit::<Fr>::new(1 << nv, gate);
121    assert!(circuit.is_satisfied());
122    let index = circuit.index;
123    //==========================================================
124    // generate pk and vks
125    let start = Instant::now();
126    for _ in 0..repetition {
127        let (_pk, _vk) = <PolyIOP<Fr> as HyperPlonkSNARK<
128            Bls12_381,
129            MultilinearKzgPCS<Bls12_381>,
130        >>::preprocess(&index, &pcs_srs)?;
131    }
132    println!(
133        "key extraction for {} variables: {} us",
134        nv,
135        start.elapsed().as_micros() / repetition as u128
136    );
137    let (pk, vk) =
138        <PolyIOP<Fr> as HyperPlonkSNARK<Bls12_381, MultilinearKzgPCS<Bls12_381>>>::preprocess(
139            &index, &pcs_srs,
140        )?;
141    //==========================================================
142    // generate a proof
143    let start = Instant::now();
144    for _ in 0..repetition {
145        let _proof =
146            <PolyIOP<Fr> as HyperPlonkSNARK<Bls12_381, MultilinearKzgPCS<Bls12_381>>>::prove(
147                &pk,
148                &circuit.public_inputs,
149                &circuit.witnesses,
150            )?;
151    }
152    let t = start.elapsed().as_micros() / repetition as u128;
153    println!(
154        "proving for {} variables: {} us",
155        nv,
156        start.elapsed().as_micros() / repetition as u128
157    );
158    file.write_all(format!("{} {}\n", nv, t).as_ref()).unwrap();
159
160    let proof = <PolyIOP<Fr> as HyperPlonkSNARK<Bls12_381, MultilinearKzgPCS<Bls12_381>>>::prove(
161        &pk,
162        &circuit.public_inputs,
163        &circuit.witnesses,
164    )?;
165    //==========================================================
166    // verify a proof
167    let start = Instant::now();
168    for _ in 0..repetition {
169        let verify =
170            <PolyIOP<Fr> as HyperPlonkSNARK<Bls12_381, MultilinearKzgPCS<Bls12_381>>>::verify(
171                &vk,
172                &circuit.public_inputs,
173                &proof,
174            )?;
175        assert!(verify);
176    }
177    println!(
178        "verifying for {} variables: {} us",
179        nv,
180        start.elapsed().as_micros() / repetition as u128
181    );
182    Ok(())
183}