midnight-proofs 0.7.1

use std::{collections::HashMap, iter};

use ff::{PrimeField, WithSmallOrderMulGroup};
use rand_chacha::ChaCha20Rng;
use rand_core::{CryptoRng, RngCore, SeedableRng};
use rayon::current_num_threads;

use super::Argument;
use crate::{
    plonk::Error,
    poly::{
        commitment::PolynomialCommitmentScheme, Coeff, EvaluationDomain, ExtendedLagrangeCoeff,
        Polynomial, ProverQuery,
    },
    transcript::{Hashable, Transcript},
    utils::arithmetic::{eval_polynomial, parallelize},
};

#[cfg_attr(feature = "bench-internal", derive(Clone))]
#[derive(Debug)]
pub(crate) struct Committed<F: PrimeField> {
    pub(crate) random_poly: Polynomial<F, Coeff>,
}

#[cfg_attr(feature = "bench-internal", derive(Clone))]
pub(crate) struct Constructed<F: PrimeField> {
    h_pieces: Vec<Polynomial<F, Coeff>>,
    committed: Committed<F>,
}

pub(crate) struct Evaluated<F: PrimeField> {
    h_poly: Polynomial<F, Coeff>,
    committed: Committed<F>,
}

impl<F: WithSmallOrderMulGroup<3>, CS: PolynomialCommitmentScheme<F>> Argument<F, CS> {
    pub(crate) fn commit<T: Transcript>(
        params: &CS::Parameters,
        domain: &EvaluationDomain<F>,
        rng: &mut (impl RngCore + CryptoRng),
        transcript: &mut T,
    ) -> Result<Committed<F>, Error>
    where
        CS::Commitment: Hashable<T::Hash>,
        F: Hashable<T::Hash>,
    {
        // Sample a random polynomial of degree n - 1
        let n = 1usize << domain.k() as usize;
        let mut rand_vec = vec![F::ZERO; n];

        let num_threads = current_num_threads();
        let chunk_size = n / num_threads;
        let thread_seeds = (0..)
            .step_by(chunk_size + 1)
            .take(n % num_threads)
            .chain(
                (chunk_size != 0)
                    .then(|| ((n % num_threads) * (chunk_size + 1)..).step_by(chunk_size))
                    .into_iter()
                    .flatten(),
            )
            .take(num_threads)
            .zip(iter::repeat_with(|| {
                let mut seed = [0u8; 32];
                rng.fill_bytes(&mut seed);
                ChaCha20Rng::from_seed(seed)
            }))
            .collect::<HashMap<_, _>>();

        parallelize(&mut rand_vec, |chunk, offset| {
            let mut rng = thread_seeds[&offset].clone();
            chunk.iter_mut().for_each(|v| *v = F::random(&mut rng));
        });

        let random_poly: Polynomial<F, Coeff> = domain.coeff_from_vec(rand_vec);

        // Commit
        let c = CS::commit(params, &random_poly);
        transcript.write(&c)?;

        Ok(Committed { random_poly })
    }
}

impl<F: WithSmallOrderMulGroup<3>> Committed<F> {
    pub(crate) fn construct<CS: PolynomialCommitmentScheme<F>, T: Transcript>(
        self,
        params: &CS::Parameters,
        domain: &EvaluationDomain<F>,
        h_poly: Polynomial<F, ExtendedLagrangeCoeff>,
        rng: &mut (impl RngCore + CryptoRng),
        transcript: &mut T,
    ) -> Result<Constructed<F>, Error>
    where
        CS::Commitment: Hashable<T::Hash>,
        F: Hashable<T::Hash>,
    {
        // Divide by t(X) = X^{params.n} - 1.
        let h_poly = domain.divide_by_vanishing_poly(h_poly);

        // Obtain final h(X) polynomial
        let mut h_poly = domain.extended_to_coeff(h_poly);

        // Let n := size of evaluation domain
        // Let d := degree of constraint system
        // Hence, the degree of the quotient poly is: d*(n-1) - n = (d-1)*(n-1) - 1,
        // and a domain of size (d-1)*(n-1) suffices to correctly represent it
        h_poly.truncate((domain.n - 1) as usize * domain.get_quotient_poly_degree());

        // Split h(X) up into pieces
        let mut h_pieces = h_poly
            .chunks_exact((domain.n - 1) as usize)
            .map(|v| v.to_vec())
            .collect::<Vec<_>>();
        drop(h_poly);

        blind_quotient_limbs(&mut h_pieces, rng);

        let h_pieces: Vec<_> =
            h_pieces.into_iter().map(|h_piece| domain.coeff_from_vec(h_piece)).collect();

        // Compute commitments to each h(X) piece
        let h_commitments: Vec<_> =
            h_pieces.iter().map(|h_piece| CS::commit(params, h_piece)).collect();

        // Hash each h(X) piece
        for c in h_commitments {
            transcript.write(&c)?;
        }

        Ok(Constructed {
            h_pieces,
            committed: self,
        })
    }
}

fn blind_quotient_limbs<F: PrimeField>(
    quotient_limbs: &mut [Vec<F>],
    rng: &mut (impl RngCore + CryptoRng),
) {
    let nr_limbs = quotient_limbs.len();
    assert!(nr_limbs >= 2);

    for i in 1..nr_limbs {
        let t = F::random(&mut *rng);
        quotient_limbs[i - 1].push(t);
        quotient_limbs[i][0] -= t;
    }

    quotient_limbs[nr_limbs - 1].push(F::ZERO);
}

impl<F: WithSmallOrderMulGroup<3>> Constructed<F> {
    pub(crate) fn evaluate<T: Transcript>(
        self,
        x: F,
        domain: &EvaluationDomain<F>,
        transcript: &mut T,
    ) -> Result<Evaluated<F>, Error>
    where
        F: Hashable<T::Hash>,
    {
        let splitting_factor: F = x.pow_vartime([domain.n - 1]);
        let h_poly = self
            .h_pieces
            .into_iter()
            .rev()
            .reduce(|acc, eval| acc * splitting_factor + eval)
            .expect("H pieces should not be empty");

        let random_eval = eval_polynomial(&self.committed.random_poly, x);
        transcript.write(&random_eval)?;

        Ok(Evaluated {
            h_poly,
            committed: self.committed,
        })
    }
}

impl<F: PrimeField> Evaluated<F> {
    pub(crate) fn open(&self, x: F) -> impl Iterator<Item = ProverQuery<'_, F>> + Clone {
        iter::empty()
            .chain(Some(ProverQuery {
                point: x,
                poly: &self.h_poly,
            }))
            .chain(Some(ProverQuery {
                point: x,
                poly: &self.committed.random_poly,
            }))
    }
}