plonky2/evm/src/stark_testing.rs

use anyhow::{ensure, Result};
use plonky2::field::extension::{Extendable, FieldExtension};
use plonky2::field::polynomial::{PolynomialCoeffs, PolynomialValues};
use plonky2::field::types::{Field, Sample};
use plonky2::hash::hash_types::RichField;
use plonky2::iop::witness::{PartialWitness, WitnessWrite};
use plonky2::plonk::circuit_builder::CircuitBuilder;
use plonky2::plonk::circuit_data::CircuitConfig;
use plonky2::plonk::config::GenericConfig;
use plonky2::util::transpose;
use plonky2_util::{log2_ceil, log2_strict};

use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
use crate::evaluation_frame::StarkEvaluationFrame;
use crate::stark::Stark;

const WITNESS_SIZE: usize = 1 << 5;

/// Tests that the constraints imposed by the given STARK are low-degree by applying them to random
/// low-degree witness polynomials.
pub(crate) fn test_stark_low_degree<
    F: RichField + Extendable<D>,
    S: Stark<F, D>,
    const D: usize,
>(
    stark: S,
) -> Result<()> {
    let rate_bits = log2_ceil(stark.constraint_degree() + 1);

    let trace_ldes = random_low_degree_matrix::<F>(S::COLUMNS, rate_bits);
    let size = trace_ldes.len();

    let lagrange_first = PolynomialValues::selector(WITNESS_SIZE, 0).lde(rate_bits);
    let lagrange_last = PolynomialValues::selector(WITNESS_SIZE, WITNESS_SIZE - 1).lde(rate_bits);

    let last = F::primitive_root_of_unity(log2_strict(WITNESS_SIZE)).inverse();
    let subgroup =
        F::cyclic_subgroup_known_order(F::primitive_root_of_unity(log2_strict(size)), size);
    let alpha = F::rand();
    let constraint_evals = (0..size)
        .map(|i| {
            let vars = S::EvaluationFrame::from_values(
                &trace_ldes[i],
                &trace_ldes[(i + (1 << rate_bits)) % size],
            );

            let mut consumer = ConstraintConsumer::<F>::new(
                vec![alpha],
                subgroup[i] - last,
                lagrange_first.values[i],
                lagrange_last.values[i],
            );
            stark.eval_packed_base(&vars, &mut consumer);
            consumer.accumulators()[0]
        })
        .collect::<Vec<_>>();

    let constraint_poly_values = PolynomialValues::new(constraint_evals);
    if !constraint_poly_values.is_zero() {
        let constraint_eval_degree = constraint_poly_values.degree();
        let maximum_degree = WITNESS_SIZE * stark.constraint_degree() - 1;

        ensure!(
            constraint_eval_degree <= maximum_degree,
            "Expected degrees at most {} * {} - 1 = {}, actual {:?}",
            WITNESS_SIZE,
            stark.constraint_degree(),
            maximum_degree,
            constraint_eval_degree
        );
    }

    Ok(())
}

/// Tests that the circuit constraints imposed by the given STARK are coherent with the native constraints.
pub(crate) fn test_stark_circuit_constraints<
    F: RichField + Extendable<D>,
    C: GenericConfig<D, F = F>,
    S: Stark<F, D>,
    const D: usize,
>(
    stark: S,
) -> Result<()> {
    // Compute native constraint evaluation on random values.
    let vars = S::EvaluationFrame::from_values(
        &F::Extension::rand_vec(S::COLUMNS),
        &F::Extension::rand_vec(S::COLUMNS),
    );

    let alphas = F::rand_vec(1);
    let z_last = F::Extension::rand();
    let lagrange_first = F::Extension::rand();
    let lagrange_last = F::Extension::rand();
    let mut consumer = ConstraintConsumer::<F::Extension>::new(
        alphas
            .iter()
            .copied()
            .map(F::Extension::from_basefield)
            .collect(),
        z_last,
        lagrange_first,
        lagrange_last,
    );
    stark.eval_ext(&vars, &mut consumer);
    let native_eval = consumer.accumulators()[0];

    // Compute circuit constraint evaluation on same random values.
    let circuit_config = CircuitConfig::standard_recursion_config();
    let mut builder = CircuitBuilder::<F, D>::new(circuit_config);
    let mut pw = PartialWitness::<F>::new();

    let locals_t = builder.add_virtual_extension_targets(S::COLUMNS);
    pw.set_extension_targets(&locals_t, vars.get_local_values());
    let nexts_t = builder.add_virtual_extension_targets(S::COLUMNS);
    pw.set_extension_targets(&nexts_t, vars.get_next_values());
    let alphas_t = builder.add_virtual_targets(1);
    pw.set_target(alphas_t[0], alphas[0]);
    let z_last_t = builder.add_virtual_extension_target();
    pw.set_extension_target(z_last_t, z_last);
    let lagrange_first_t = builder.add_virtual_extension_target();
    pw.set_extension_target(lagrange_first_t, lagrange_first);
    let lagrange_last_t = builder.add_virtual_extension_target();
    pw.set_extension_target(lagrange_last_t, lagrange_last);

    let vars = S::EvaluationFrameTarget::from_values(&locals_t, &nexts_t);
    let mut consumer = RecursiveConstraintConsumer::<F, D>::new(
        builder.zero_extension(),
        alphas_t,
        z_last_t,
        lagrange_first_t,
        lagrange_last_t,
    );
    stark.eval_ext_circuit(&mut builder, &vars, &mut consumer);
    let circuit_eval = consumer.accumulators()[0];
    let native_eval_t = builder.constant_extension(native_eval);
    builder.connect_extension(circuit_eval, native_eval_t);

    let data = builder.build::<C>();
    let proof = data.prove(pw)?;
    data.verify(proof)
}

fn random_low_degree_matrix<F: Field>(num_polys: usize, rate_bits: usize) -> Vec<Vec<F>> {
    let polys = (0..num_polys)
        .map(|_| random_low_degree_values(rate_bits))
        .collect::<Vec<_>>();

    transpose(&polys)
}

fn random_low_degree_values<F: Field>(rate_bits: usize) -> Vec<F> {
    PolynomialCoeffs::new(F::rand_vec(WITNESS_SIZE))
        .lde(rate_bits)
        .fft()
        .values
}