p3-commit 0.8.0

mod pcs;

use alloc::vec::Vec;
use core::marker::PhantomData;

use p3_challenger::CanSample;
use p3_dft::TwoAdicSubgroupDft;
use p3_field::coset::TwoAdicMultiplicativeCoset;
use p3_field::{ExtensionField, Field, TwoAdicField};
use p3_matrix::Matrix;
use p3_matrix::dense::RowMajorMatrix;
use p3_util::log2_strict_usize;
use p3_util::zip_eq::zip_eq;
pub use pcs::assert_pcs_opening_contract;
use serde::{Deserialize, Serialize};

use crate::{
    CommitmentOpening, MatrixOpening, OpenedValues, OpeningRequest, Pcs, PointOpening,
    PolynomialSpace, UnivariateStarkPcs,
};

/// A trivial PCS: its commitment is simply the coefficients of each poly.
#[derive(Clone, Debug)]
pub struct TrivialPcs<Val: TwoAdicField, Dft: TwoAdicSubgroupDft<Val>> {
    pub dft: Dft,
    // degree bound
    pub log_n: usize,
    pub _phantom: PhantomData<Val>,
}

pub fn eval_coeffs_at_pt<F: Field, EF: ExtensionField<F>>(
    coeffs: &RowMajorMatrix<F>,
    x: EF,
) -> Vec<EF> {
    let mut acc = EF::zero_vec(coeffs.width());
    for r in (0..coeffs.height()).rev() {
        let row = coeffs.row_slice(r).unwrap();
        for (acc_c, row_c) in acc.iter_mut().zip(row.iter()) {
            *acc_c *= x;
            *acc_c += *row_c;
        }
    }
    acc
}

impl<Val, Dft, Challenge, Challenger> Pcs<Challenge, Challenger> for TrivialPcs<Val, Dft>
where
    Val: TwoAdicField,
    Challenge: ExtensionField<Val>,
    Challenger: CanSample<Challenge>,
    Dft: TwoAdicSubgroupDft<Val>,
    Vec<Vec<Val>>: Serialize + for<'de> Deserialize<'de>,
{
    type Domain = TwoAdicMultiplicativeCoset<Val>;
    type Commitment = Vec<Vec<Val>>;
    type ProverData = Vec<RowMajorMatrix<Val>>;
    type Proof = ();
    type Error = ();
    type ProverError = core::convert::Infallible;

    fn natural_domain_for_degree(&self, degree: usize) -> Self::Domain {
        // This panics if (and only if) `degree` is not a power of 2 or `degree`
        // > `1 << Val::TWO_ADICITY`.
        TwoAdicMultiplicativeCoset::new(Val::ONE, log2_strict_usize(degree)).unwrap()
    }

    fn commit(
        &self,
        evaluations: impl IntoIterator<Item = (Self::Domain, RowMajorMatrix<Val>)>,
    ) -> Result<(Self::Commitment, Self::ProverData), Self::ProverError> {
        let coeffs: Vec<_> = evaluations
            .into_iter()
            .map(|(domain, evals)| {
                let log_domain_size = log2_strict_usize(domain.size());
                // for now, only commit on larger domain than natural
                assert!(log_domain_size >= self.log_n);
                assert_eq!(domain.size(), evals.height());
                // coset_idft_batch
                let mut coeffs = self.dft.idft_batch(evals);
                coeffs
                    .rows_mut()
                    .zip(domain.shift_inverse().powers())
                    .for_each(|(row, weight)| {
                        row.iter_mut().for_each(|coeff| {
                            *coeff *= weight;
                        });
                    });
                coeffs
            })
            .collect();
        Ok((
            coeffs.clone().into_iter().map(|m| m.values).collect(),
            coeffs,
        ))
    }

    fn open(
        &self,
        // For each round,
        rounds: Vec<OpeningRequest<'_, Self::ProverData, Challenge>>,
        _challenger: &mut Challenger,
    ) -> Result<(OpenedValues<Challenge>, Self::Proof), Self::ProverError> {
        Ok((
            rounds
                .into_iter()
                .map(
                    |OpeningRequest {
                         prover_data: coeffs_for_round,
                         points: points_for_round,
                     }| {
                        // ensure that each matrix corresponds to a set of opening points
                        debug_assert_eq!(coeffs_for_round.len(), points_for_round.len());
                        coeffs_for_round
                            .iter()
                            .zip(points_for_round)
                            .map(|(coeffs_for_mat, points_for_mat)| {
                                points_for_mat
                                    .into_iter()
                                    .map(|pt| eval_coeffs_at_pt(coeffs_for_mat, pt))
                                    .collect()
                            })
                            .collect()
                    },
                )
                .collect(),
            (),
        ))
    }

    // This is a testing function, so we allow panics for convenience.
    #[allow(clippy::panic_in_result_fn)]
    fn verify(
        &self,
        // For each round:
        rounds: Vec<CommitmentOpening<Challenge, Self::Commitment, Self::Domain>>,
        _proof: &Self::Proof,
        _challenger: &mut Challenger,
    ) -> Result<(), Self::Error> {
        for CommitmentOpening {
            commitment: comm,
            matrices: round_opening,
        } in rounds
        {
            for (
                coeff_vec,
                MatrixOpening {
                    domain,
                    points: points_and_values,
                },
            ) in zip_eq(comm, round_opening, ())?
            {
                let width = coeff_vec.len() / domain.size();
                assert_eq!(width * domain.size(), coeff_vec.len());
                let coeffs = RowMajorMatrix::new(coeff_vec, width);
                for PointOpening { point: pt, values } in points_and_values {
                    assert_eq!(eval_coeffs_at_pt(&coeffs, pt), values);
                }
            }
        }
        Ok(())
    }
}

impl<Val, Dft, Challenge, Challenger> UnivariateStarkPcs<Challenge, Challenger>
    for TrivialPcs<Val, Dft>
where
    Val: TwoAdicField,
    Challenge: ExtensionField<Val>,
    Challenger: CanSample<Challenge>,
    Dft: TwoAdicSubgroupDft<Val>,
    Vec<Vec<Val>>: Serialize + for<'de> Deserialize<'de>,
{
    type EvaluationsOnDomain<'a> = Dft::Evaluations;

    const ZK: bool = false;

    fn log_max_trace_height(&self) -> usize {
        Val::TWO_ADICITY
    }

    fn log_min_trace_height(&self) -> usize {
        // Multiplicative coset selectors are defined at every height, down to a single row.
        0
    }

    fn commit_quotient(
        &self,
        quotient_domain: Self::Domain,
        quotient_evaluations: RowMajorMatrix<crate::Val<Self::Domain>>,
        num_chunks: usize,
    ) -> Result<(Self::Commitment, Self::ProverData), Self::ProverError> {
        let quotient_sub_evaluations =
            quotient_domain.split_evals(num_chunks, quotient_evaluations);
        let quotient_sub_domains = quotient_domain.split_domains(num_chunks);

        Pcs::<Challenge, Challenger>::commit(
            self,
            quotient_sub_domains
                .into_iter()
                .zip(quotient_sub_evaluations),
        )
    }

    fn get_quotient_ldes(
        &self,
        _evaluations: impl IntoIterator<Item = (Self::Domain, RowMajorMatrix<Val>)>,
        _num_chunks: usize,
    ) -> Result<Vec<RowMajorMatrix<crate::Val<Self::Domain>>>, Self::ProverError> {
        unimplemented!("This PCS does not support computing of LDEs");
    }

    fn commit_ldes(
        &self,
        _ldes: Vec<RowMajorMatrix<Val>>,
    ) -> Result<(Self::Commitment, Self::ProverData), Self::ProverError> {
        unimplemented!("This PCS does not support computing of LDEs");
    }

    fn get_evaluations_on_domain<'a>(
        &self,
        prover_data: &'a Self::ProverData,
        idx: usize,
        domain: Self::Domain,
    ) -> Self::EvaluationsOnDomain<'a> {
        let mut coeffs = prover_data[idx].clone();
        assert!(domain.log_size() >= self.log_n);
        coeffs.values.resize(
            coeffs.values.len() << (domain.log_size() - self.log_n),
            Val::ZERO,
        );
        self.dft.coset_dft_batch(coeffs, domain.shift())
    }
}