provekit-whir 0.1.1

//! Example: whir (non-ZK) batch commitment for 44 polynomials of size 2^19.
//!
//! Run with:
//! ```
//! RUSTFLAGS="-C target-cpu=native" cargo run --release --features asm --bin batch_commit
//! ```

use std::{borrow::Cow, time::Instant};

use ark_ff::FftField;
use ark_std::rand::{distributions::Standard, prelude::Distribution, thread_rng, Rng};
use provekit_whir::{
    algebra::{
        embedding::{Embedding, Identity},
        fields::Field256,
        linear_form::{Evaluate, LinearForm, MultilinearExtension},
    },
    hash::HASH_COUNTER,
    parameters::ProtocolParameters,
    transcript::{codecs::Empty, Codec, DomainSeparator, ProverState, VerifierState},
};

const NUM_POLYNOMIALS: usize = 42;
const NUM_VARIABLES: usize = 19; // polynomial size = 2^19
const NUM_COEFFS: usize = 1 << NUM_VARIABLES;
const NUM_EVAL_POINTS: usize = 1;

fn main() {
    run_batch::<Identity<Field256>>();
}

fn run_batch<M>()
where
    M: Embedding + Default,
    M::Source: FftField,
    M::Target: FftField + Codec,
    Standard: Distribution<M::Source>,
    Standard: Distribution<M::Target>,
{
    use provekit_whir::protocols::whir::Config;

    let mut rng = thread_rng();

    let security_level: usize = 128;
    let pow_bits = 10;
    let starting_rate = 1;
    let initial_folding_factor = 4;
    let folding_factor = 4;

    let whir_params = ProtocolParameters {
        security_level,
        pow_bits,
        initial_folding_factor,
        folding_factor,
        unique_decoding: false,
        starting_log_inv_rate: starting_rate,
        batch_size: NUM_POLYNOMIALS,
        hash_id: provekit_whir::hash::BLAKE3,
    };

    let params = Config::<M>::new(1 << NUM_VARIABLES, &whir_params);

    println!("=========================================");
    println!("whir Batch Commit Example (non-ZK)");
    println!("=========================================");
    println!("Polynomials:     {NUM_POLYNOMIALS}");
    println!("Variables:       {NUM_VARIABLES} (size = {NUM_COEFFS})");
    println!("Eval points:     {NUM_EVAL_POINTS}");
    println!("Security level:  {security_level}");
    println!("PoW bits:        {pow_bits}");
    println!("Hash:            Blake3");
    println!("{params}");

    // Generate random polynomials.
    let vectors: Vec<Vec<M::Source>> = (0..NUM_POLYNOMIALS)
        .map(|_| (0..NUM_COEFFS).map(|_| rng.gen()).collect())
        .collect();
    let vec_refs: Vec<&[M::Source]> = vectors.iter().map(Vec::as_slice).collect();

    // Build evaluation constraints.
    let points: Vec<Vec<M::Target>> = (0..NUM_EVAL_POINTS)
        .map(|_| (0..NUM_VARIABLES).map(|_| rng.gen()).collect())
        .collect();

    let forms: Vec<MultilinearExtension<M::Target>> = points
        .iter()
        .map(|p| MultilinearExtension::new(p.clone()))
        .collect();

    // Compute evaluations: row-major [form_0(poly_0), ..., form_k(poly_n-1)].
    let evaluations: Vec<M::Target> = forms
        .iter()
        .flat_map(|form| {
            vec_refs
                .iter()
                .map(|v| form.evaluate(params.embedding(), v))
        })
        .collect();

    // Convert to trait objects.
    let prove_forms: Vec<Box<dyn LinearForm<M::Target>>> = forms
        .iter()
        .map(|f| Box::new(f.clone()) as Box<dyn LinearForm<M::Target>>)
        .collect();

    // Set up transcript.
    let ds = DomainSeparator::protocol(&params)
        .session(&format!("batch_commit_44x2^19 at {}:{}", file!(), line!()))
        .instance(&Empty);

    let mut prover_state = ProverState::new_std(&ds);

    // --- Commit ---
    println!("-----------------------------------------");
    println!("Committing to {NUM_POLYNOMIALS} polynomials...");
    let commit_time = Instant::now();
    let witness = params.commit(&mut prover_state, &vec_refs);
    let commit_time = commit_time.elapsed();
    println!("Commit time:     {commit_time:.2?}");

    // --- Prove ---
    println!("Proving...");
    let prove_time = Instant::now();
    let _ = params.prove(
        &mut prover_state,
        vectors
            .iter()
            .map(|v| Cow::Borrowed(v.as_slice()))
            .collect(),
        vec![Cow::Owned(witness)],
        prove_forms,
        Cow::Borrowed(&evaluations),
    );
    let prove_time = prove_time.elapsed();
    println!("Prove time:      {prove_time:.2?}");

    let proof = prover_state.proof();
    let proof_size = proof.narg_string.len() + proof.hints.len();
    println!(
        "Proof size:      {:.1} KiB ({} bytes)",
        proof_size as f64 / 1024.0,
        proof_size,
    );
    println!("Total prover:    {:.2?}", commit_time + prove_time);

    // --- Verify ---
    let verify_forms: Vec<Box<dyn Evaluate<M>>> = forms
        .iter()
        .map(|f| Box::new(f.clone()) as Box<dyn Evaluate<M>>)
        .collect();
    let verify_form_refs: Vec<&dyn LinearForm<M::Target>> = verify_forms
        .iter()
        .map(|f| f.as_ref() as &dyn LinearForm<M::Target>)
        .collect();

    HASH_COUNTER.reset();
    let reps = 10;
    println!("-----------------------------------------");
    println!("Verifying ({reps} repetitions)...");
    let verify_time = Instant::now();
    for _ in 0..reps {
        let mut verifier_state = VerifierState::new_std(&ds, &proof);
        let commitment = params
            .receive_commitment(&mut verifier_state)
            .expect("receive commitment");
        let final_claim = params
            .verify(&mut verifier_state, &[&commitment], &evaluations)
            .expect("verification failed");
        final_claim
            .verify(verify_form_refs.iter().copied())
            .expect("final claim check failed");
    }
    let verify_elapsed = verify_time.elapsed();
    println!(
        "Verify time:     {:.2?} (avg over {reps})",
        verify_elapsed / reps as u32,
    );
    println!(
        "Avg hashes:      {:.1}k",
        (HASH_COUNTER.get() as f64 / reps as f64) / 1000.0,
    );
    println!("=========================================");
    println!("All verifications passed.");
}