// Copyright (c) Facebook, Inc. and its affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use crypto::{BatchMerkleProof, ElementHasher, Hasher};
use math::FieldElement;
use utils::{
    collections::Vec, string::ToString, ByteReader, ByteWriter, Deserializable,
    DeserializationError, Serializable, SliceReader,
};

// FRI PROOF
// ================================================================================================

/// A proof generated by a FRI prover.
///
/// A FRI proof contains information proving that a function *f* is a polynomial of some bounded
/// degree *d*. FRI proofs cannot be instantiated directly - they must be generated by a instance
/// of a [FriProver](crate::FriProver), and can be verified by a instance of a
/// [FriVerifier](crate::FriVerifier) via [VerifierChannel](crate::VerifierChannel) interface.
///
/// A proof consists of zero or more layers and a remainder polynomial. Each layer contains a set of
/// polynomial evaluations at positions queried by the verifier as well as Merkle authentication
/// paths for these evaluations (the Merkle paths are compressed into a batch Merkle proof). The
/// remainder polynomial is given by its list of coefficients i.e. field elements.
///
/// All values in a proof are stored as vectors of bytes. Thus, the values must be parsed before
/// they can be returned to the user. To do this, [parse_layers()](FriProof::parse_layers())
/// and [parse_remainder()](FriProof::parse_remainder()) methods can be used.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct FriProof {
    layers: Vec<FriProofLayer>,
    remainder: Vec<u8>,
    num_partitions: u8, // stored as power of 2
}

impl FriProof {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Creates a new FRI proof from the provided layers and remainder polynomial.
    ///
    /// # Panics
    /// Panics if:
    /// * Number of remainder elements zero or is not a power of two.
    /// * `num_partitions` is zero or is not a power of two.
    pub(crate) fn new<E: FieldElement>(
        layers: Vec<FriProofLayer>,
        remainder: Vec<E>,
        num_partitions: usize,
    ) -> Self {
        assert!(
            !remainder.is_empty(),
            "number of remainder elements must be greater than zero"
        );
        assert!(
            remainder.len().is_power_of_two(),
            "size of the remainder must be a power of two, but was {}",
            remainder.len()
        );
        assert!(
            num_partitions > 0,
            "number of partitions must be greater than zero"
        );
        assert!(
            num_partitions.is_power_of_two(),
            "number of partitions must be a power of two, but was {num_partitions}"
        );
        FriProof {
            layers,
            remainder: remainder.to_bytes(),
            num_partitions: num_partitions.trailing_zeros() as u8,
        }
    }

    // PUBLIC ACCESSORS
    // --------------------------------------------------------------------------------------------

    /// Returns the number of layers in this proof.
    pub fn num_layers(&self) -> usize {
        self.layers.len()
    }

    /// Returns the number of remainder elements in this proof.
    ///
    /// The number of elements is computed by dividing the number of remainder bytes by the size
    /// of the field element specified by `E` type parameter.
    pub fn num_remainder_elements<E: FieldElement>(&self) -> usize {
        self.remainder.len() / E::ELEMENT_BYTES
    }

    /// Returns the number of partitions used during proof generation.
    pub fn num_partitions(&self) -> usize {
        2usize.pow(self.num_partitions as u32)
    }

    /// Returns the size of this proof in bytes.
    pub fn size(&self) -> usize {
        // +1 for number of layers, +1 for remainder length, +1 for number of partitions
        self.layers
            .iter()
            .fold(self.remainder.len() + 3, |acc, layer| acc + layer.size())
    }

    // PARSING
    // --------------------------------------------------------------------------------------------

    /// Decomposes this proof into vectors of query values for each layer and corresponding Merkle
    /// authentication paths for each query (grouped into batch Merkle proofs).
    ///
    /// # Panics
    /// Panics if:
    /// * `domain_size` is not a power of two.
    /// * `folding_factor` is smaller than two or is not a power of two.
    ///
    /// # Errors
    /// Returns an error if:
    /// * This proof is not consistent with the specified `domain_size` and `folding_factor`.
    /// * Any of the layers could not be parsed successfully.
    #[allow(clippy::type_complexity)]
    pub fn parse_layers<H, E>(
        self,
        mut domain_size: usize,
        folding_factor: usize,
    ) -> Result<(Vec<Vec<E>>, Vec<BatchMerkleProof<H>>), DeserializationError>
    where
        E: FieldElement,
        H: ElementHasher<BaseField = E::BaseField>,
    {
        assert!(
            domain_size.is_power_of_two(),
            "domain size must be a power of two"
        );
        assert!(
            folding_factor.is_power_of_two(),
            "folding factor must be a power of two"
        );
        assert!(folding_factor > 1, "folding factor must be greater than 1");

        let mut layer_proofs = Vec::new();
        let mut layer_queries = Vec::new();

        // parse all layers
        for (i, layer) in self.layers.into_iter().enumerate() {
            domain_size /= folding_factor;
            let (qv, mp) = layer.parse(domain_size, folding_factor).map_err(|err| {
                DeserializationError::InvalidValue(format!("failed to parse FRI layer {i}: {err}"))
            })?;
            layer_proofs.push(mp);
            layer_queries.push(qv);
        }

        Ok((layer_queries, layer_proofs))
    }

    /// Returns a vector of remainder values (last FRI layer) parsed from this proof.
    ///
    /// # Errors
    /// Returns an error if:
    /// * The number of remainder values implied by a combination of `E` type parameter and
    ///   the number of remainder bytes in this proof is not a power of two.
    /// * Any of the remainder values could not be parsed correctly.
    /// * Not all bytes have been consumed while parsing remainder values.
    pub fn parse_remainder<E: FieldElement>(&self) -> Result<Vec<E>, DeserializationError> {
        let num_elements = self.num_remainder_elements::<E>();
        if !num_elements.is_power_of_two() {
            return Err(DeserializationError::InvalidValue(format!(
                "number of remainder values must be a power of two, but {num_elements} was implied"
            )));
        }
        let mut reader = SliceReader::new(&self.remainder);
        let remainder = E::read_batch_from(&mut reader, num_elements).map_err(|err| {
            DeserializationError::InvalidValue(format!("failed to parse FRI remainder: {err}"))
        })?;
        if reader.has_more_bytes() {
            return Err(DeserializationError::UnconsumedBytes);
        }
        Ok(remainder)
    }
}

// SERIALIZATION / DESERIALIZATION
// ------------------------------------------------------------------------------------------------

impl Serializable for FriProof {
    /// Serializes `self` and writes the resulting bytes into the `target` writer.
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        // write layers
        target.write_u8(self.layers.len() as u8);
        for layer in self.layers.iter() {
            layer.write_into(target);
        }

        // write remainder
        target.write_u16(self.remainder.len() as u16);
        target.write_bytes(&self.remainder);

        // write number of partitions
        target.write_u8(self.num_partitions);
    }
}

impl Deserializable for FriProof {
    /// Reads a FRI proof from the specified `source` and returns the result.
    ///
    /// # Errors
    /// Returns an error if a valid proof could not be read from the source.
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        // read layers
        let num_layers = source.read_u8()? as usize;
        let layers = FriProofLayer::read_batch_from(source, num_layers)?;

        // read remainder
        let num_remainder_bytes = source.read_u16()? as usize;
        let remainder = source.read_vec(num_remainder_bytes)?;

        // read number of partitions
        let num_partitions = source.read_u8()?;

        Ok(FriProof {
            layers,
            remainder,
            num_partitions,
        })
    }
}

// FRI PROOF LAYER
// ================================================================================================

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct FriProofLayer {
    values: Vec<u8>,
    paths: Vec<u8>,
}

impl FriProofLayer {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Creates a new proof layer from the specified query values and the corresponding Merkle
    /// paths aggregated into a single batch Merkle proof.
    ///
    /// # Panics
    /// Panics if `query_values` is an empty slice.
    pub(crate) fn new<H: Hasher, E: FieldElement, const N: usize>(
        query_values: Vec<[E; N]>,
        merkle_proof: BatchMerkleProof<H>,
    ) -> Self {
        assert!(!query_values.is_empty(), "query values cannot be empty");

        // TODO: add debug check that values actually hash into the leaf nodes of the batch proof

        // concatenate all query values and all internal Merkle proof nodes into vectors of bytes;
        // we care about internal nodes only because leaf nodes can be reconstructed from hashes
        // of query values
        FriProofLayer {
            values: query_values.to_bytes(),
            paths: merkle_proof.serialize_nodes(),
        }
    }

    // PUBLIC ACCESSORS
    // --------------------------------------------------------------------------------------------

    /// Returns the size of this proof layer in bytes.
    pub fn size(&self) -> usize {
        // +4 for length of values, +4 for length of paths
        self.values.len() + 4 + self.paths.len() + 4
    }

    // PARSING
    // --------------------------------------------------------------------------------------------
    /// Decomposes this layer into a combination of query values and corresponding Merkle
    /// authentication paths (grouped together into a single batch Merkle proof).
    ///
    /// # Errors
    /// Returns an error if:
    /// * This layer does not contain at least one query.
    /// * Parsing of any of the query values or the corresponding Merkle paths fails.
    /// * Not all bytes have been consumed while parsing this layer.
    pub fn parse<H, E>(
        self,
        domain_size: usize,
        folding_factor: usize,
    ) -> Result<(Vec<E>, BatchMerkleProof<H>), DeserializationError>
    where
        E: FieldElement,
        H: ElementHasher<BaseField = E::BaseField>,
    {
        // make sure the number of value bytes can be parsed into a whole number of queries
        let num_query_bytes = E::ELEMENT_BYTES * folding_factor;
        if self.values.len() % num_query_bytes != 0 {
            return Err(DeserializationError::InvalidValue(format!(
                "number of value bytes ({}) does not divide into whole number of queries",
                self.values.len(),
            )));
        }

        let num_queries = self.values.len() / num_query_bytes;
        if num_queries == 0 {
            return Err(DeserializationError::InvalidValue(
                "a FRI layer must contain at least one query".to_string(),
            ));
        }
        let mut hashed_queries = vec![H::Digest::default(); num_queries];
        let mut query_values = Vec::with_capacity(num_queries * folding_factor);

        // read bytes corresponding to each query, convert them into field elements,
        // and also hash them to build leaf nodes of the batch Merkle proof
        let mut reader = SliceReader::new(&self.values);
        for query_hash in hashed_queries.iter_mut() {
            let mut qe = E::read_batch_from(&mut reader, folding_factor)?;
            *query_hash = H::hash_elements(&qe);
            query_values.append(&mut qe);
        }
        if reader.has_more_bytes() {
            return Err(DeserializationError::UnconsumedBytes);
        }

        // build batch Merkle proof
        let mut reader = SliceReader::new(&self.paths);
        let tree_depth = domain_size.ilog2() as u8;
        let merkle_proof = BatchMerkleProof::deserialize(&mut reader, hashed_queries, tree_depth)?;
        if reader.has_more_bytes() {
            return Err(DeserializationError::UnconsumedBytes);
        }

        Ok((query_values, merkle_proof))
    }
}

// SERIALIZATION / DESERIALIZATION
// ------------------------------------------------------------------------------------------------

impl Serializable for FriProofLayer {
    /// Serializes this proof layer and writes the resulting bytes to the specified `target`.
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        // write value bytes
        target.write_u32(self.values.len() as u32);
        target.write_bytes(&self.values);

        // write path bytes
        target.write_u32(self.paths.len() as u32);
        target.write_bytes(&self.paths);
    }
}

impl Deserializable for FriProofLayer {
    /// Reads a single proof layer form the `source` and returns it.
    ///
    /// # Errors
    /// Returns an error if a valid layer could not be read from the specified source.
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        // read values
        let num_value_bytes = source.read_u32()?;
        if num_value_bytes == 0 {
            return Err(DeserializationError::InvalidValue(
                "a FRI proof layer must contain at least one queried evaluation".to_string(),
            ));
        }
        let values = source.read_vec(num_value_bytes as usize)?;

        // read paths
        let num_paths_bytes = source.read_u32()?;
        let paths = source.read_vec(num_paths_bytes as usize)?;

        Ok(FriProofLayer { values, paths })
    }
}


// Copyright (c) Facebook, Inc. and its affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use crypto::ElementHasher;
use math::FieldElement;
use utils::{collections::Vec, iter_mut, uninit_vector};

#[cfg(feature = "concurrent")]
use utils::iterators::*;

/// Maps positions in the evaluation domain to indexes of commitment Merkle tree.
pub fn map_positions_to_indexes(
    positions: &[usize],
    source_domain_size: usize,
    folding_factor: usize,
    num_partitions: usize,
) -> Vec<usize> {
    // if there was only 1 partition, order of elements in the commitment tree
    // is the same as the order of elements in the evaluation domain
    if num_partitions == 1 {
        return positions.to_vec();
    }

    let target_domain_size = source_domain_size / folding_factor;
    let partition_size = target_domain_size / num_partitions;

    let mut result = Vec::new();
    for position in positions {
        let partition_idx = position % num_partitions;
        let local_idx = (position - partition_idx) / num_partitions;
        let position = partition_idx * partition_size + local_idx;
        result.push(position);
    }

    result
}

/// Hashes each of the arrays in the provided slice and returns a vector of resulting hashes.
pub fn hash_values<H, E, const N: usize>(values: &[[E; N]]) -> Vec<H::Digest>
where
    E: FieldElement,
    H: ElementHasher<BaseField = E::BaseField>,
{
    let mut result: Vec<H::Digest> = unsafe { uninit_vector(values.len()) };
    iter_mut!(result, 1024).zip(values).for_each(|(r, v)| {
        *r = H::hash_elements(v);
    });
    result
}


// Copyright (c) Facebook, Inc. and its affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

//! Contains functions for folding FRI layers.
//!
//! This module is exposed publicly primarily for benchmarking and documentation purposes. The
//! functions contained here are not intended to be used by the end-user of the crate.

#[cfg(feature = "concurrent")]
use utils::iterators::*;

use math::{
    fft::{get_inv_twiddles, serial_fft},
    get_power_series_with_offset, polynom, FieldElement, StarkField,
};
use utils::{collections::Vec, iter_mut, uninit_vector};

// DEGREE-RESPECTING PROJECTION
// ================================================================================================
/// Applies degree-respecting projection to evaluations of a polynomial.
///
/// This process reduces the  degree of the polynomial implied by the `evaluations` by `N`.
/// For example, for `N` = 4, this is equivalent to the following:
/// - Let `evaluations` contain the evaluations of polynomial *f*(x) of degree *k*.
/// - Group coefficients of *f* so that *f*(x) = a(x) + x * b(x) + x^2 * c(x) + x^3 * d(x)
/// - Compute random linear combination of polynomials a, b, c, d as
///   *f'*(x) = a + Î± * b + Î±^2 * c + Î±^3 * d, where Î± is an element drawn uniformly at random from
///   the entire field.
/// - evaluate f'(x) on a domain which consists of x^4 from the original domain (and thus is
///   1/4 the size).
///
/// However, the reduction is performed without converting the polynomials into coefficient form.
/// That is, we can go directly form evaluations to folded evaluations. For this, for each
/// evaluation in the folded domain, we need `N` evaluations in the source domain. For example,
/// for the case of `N` = 4, to compute the evaluation of *f'*(x) we need to have evaluations of
/// *f*(x) from the source domain at x^{1/4}, x^{2/4}, x^{3/4}, x.
///
/// This function expect the `evaluations` to be already in a transposed form such that all
/// evaluations needed to compute a single evaluation in the folded domain are next to each other.
///
/// The example below shows the equivalence of performing the projection via coefficient form and
/// via evaluation form for `N` = 2.
/// ```
/// # use math::{StarkField, FieldElement, fields::f128::BaseElement, get_power_series_with_offset, polynom};
/// # use rand_utils::{rand_value, rand_vector};
/// # use utils::transpose_slice;
/// # use winter_fri::folding::apply_drp;
/// // generate random alpha
/// let alpha: BaseElement = rand_value();
///
/// // degree 7 polynomial f(x)
/// let poly: Vec<BaseElement> = rand_vector(8);
///
/// // f'(x) = g(x) + alpha * h(x) such that g(x) consists of even coefficients of f(x) and h(x)
/// // consists of odd coefficients of f(x). This is equivalent to using `folding_factor = 2`.
/// let mut folded_poly = Vec::new();
/// for i in 0..poly.len() / 2 {
///     folded_poly.push(poly[2 * i] + alpha * poly[2 * i + 1]);
/// }
///
/// // build a domain of 32 elements
/// let n = 32_usize;
/// let offset = BaseElement::GENERATOR;
/// let g = BaseElement::get_root_of_unity(n.trailing_zeros());
/// let domain = get_power_series_with_offset(g, offset, n);
///
/// // build a folded domain of half the size, such that x in the source domain maps to
/// // x^2 in the folded domain.
/// let g = BaseElement::get_root_of_unity((n / 2).trailing_zeros());
/// let folded_domain = get_power_series_with_offset(g, offset.exp(2), n / 2);
///
/// // evaluate the polynomials over their respective domains
/// let evaluations = polynom::eval_many(&poly, &domain);
/// let folded_evaluations = polynom::eval_many(&folded_poly, &folded_domain);
///
/// // use DRP to perform folding without converting polynomials into coefficient form
/// let transposed_evaluations = transpose_slice::<BaseElement, 2>(&evaluations);
/// let drp_evaluations = apply_drp(&transposed_evaluations, offset, alpha);
///
/// // applying DRP should be equivalent to folding polynomials in coefficient form
/// assert_eq!(folded_evaluations, drp_evaluations);
/// ```
pub fn apply_drp<B, E, const N: usize>(values: &[[E; N]], domain_offset: B, alpha: E) -> Vec<E>
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    // build offset inverses and twiddles used during polynomial interpolation
    let inv_offsets = get_inv_offsets(values.len(), domain_offset, N);
    let inv_twiddles = get_inv_twiddles::<B>(N);
    let len_offset = E::inv((N as u64).into());

    let mut result = unsafe { uninit_vector(values.len()) };
    iter_mut!(result)
        .zip(values)
        .zip(inv_offsets)
        .for_each(|((result, values), domain_offset)| {
            // interpolate the values into a polynomial; this is similar to interpolation with
            // offset implemented in math::fft module
            let mut poly = *values;
            serial_fft(&mut poly, &inv_twiddles);

            let mut offset = len_offset;
            let domain_offset = E::from(domain_offset);
            for coeff in poly.iter_mut() {
                *coeff *= offset;
                offset *= domain_offset;
            }

            // evaluate the polynomial at alpha, and save the result
            *result = polynom::eval(&poly, alpha)
        });

    result
}

// POSITION FOLDING
// ================================================================================================
/// Maps positions in the source domain, to positions in the folded domain.
///
/// The size of the folded domain is reduced by the `folding_factor` as compared to the size of the
/// source domain. Thus, the original positions may fall outside of the folded domain. To map
/// positions in the source domain to positions in the folded domain, we simply compute
/// `new_position = old_position % folded_domain_size` and discard duplicates.
///
/// ```
/// # use winter_fri::folding::fold_positions;
/// let positions = [1, 9, 12, 20];
/// let folded_positions = fold_positions(&positions, 32, 4);
///
/// // 1 and 9 should map to 1; 12 and 20 should map to 4
/// assert_eq!(vec![1, 4], folded_positions);
/// ```
///
/// The domains we usually work with have the following form: c, c * Ï‰, c * Ï‰^2, ..., c * Ï‰^(*n*-1),
/// where *n* is the size of the domain, Ï‰ is *n*th root of unity, and c is the domain offset. Thus,
/// the procedure described above maps c * Ï‰^*i* to (c * Ï‰^*i*)^*k*, where *k* is the folding factor.
/// ```
/// # use math::{fields::f128::BaseElement, StarkField, FieldElement, get_power_series_with_offset};
/// // build a domain of size 32
/// let n = 32usize;
/// let c = BaseElement::GENERATOR;
/// let g = BaseElement::get_root_of_unity(n.trailing_zeros());
/// let source_domain = get_power_series_with_offset(g, c, n);
///
/// // build a domain of size 8 (folding_factor = 4)
/// let g = BaseElement::get_root_of_unity((n / 4).trailing_zeros());
/// let folded_domain = get_power_series_with_offset(g, c.exp(4_u32.into()), (n / 4));
///
/// // position 1 in the source domain maps to position 1 in the folded domain
/// assert_eq!(folded_domain[1], source_domain[1].exp(4));
///
/// // position 9 in the source domain also maps to position 1 in the folded domain
/// assert_eq!(folded_domain[1], source_domain[9].exp(4));
/// ```
pub fn fold_positions(
    positions: &[usize],
    source_domain_size: usize,
    folding_factor: usize,
) -> Vec<usize> {
    let target_domain_size = source_domain_size / folding_factor;

    let mut result = Vec::new();
    for position in positions {
        let position = position % target_domain_size;
        // make sure we don't record duplicated values
        if !result.contains(&position) {
            result.push(position);
        }
    }

    result
}

// HELPER FUNCTIONS
// ================================================================================================

fn get_inv_offsets<B>(domain_size: usize, domain_offset: B, folding_factor: usize) -> Vec<B>
where
    B: StarkField,
{
    let n = domain_size * folding_factor;
    let g = B::get_root_of_unity(n.ilog2());
    get_power_series_with_offset(g.inv(), domain_offset.inv(), domain_size)
}


// Copyright (c) Facebook, Inc. and its affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use core::marker::PhantomData;
use crypto::{ElementHasher, Hasher, RandomCoin};
use math::FieldElement;
use utils::collections::Vec;

// PROVER CHANNEL TRAIT
// ================================================================================================

/// Defines an interface for a channel over which a prover communicates with a verifier.
///
/// The prover uses this channel to send commitments to FRI layer polynomials to the verifier, and
/// then to draw a random value Î± from the channel after each commitment is sent. The prover then
/// uses this Î± to construct the next FRI layer.
///
/// In the interactive version of the protocol, the verifier chooses Î± uniformly at random from
/// the entire field. In the non-interactive version, the Î± is drawn pseudo-randomly based on the
/// commitments the prover has written into the channel up to this point.
pub trait ProverChannel<E: FieldElement> {
    /// Hash function used by the prover to commit to polynomial evaluations.
    type Hasher: Hasher;

    /// Sends a layer commitment to the verifier.
    ///
    /// A layer commitment is a root of a Merkle tree built from evaluations of a polynomial
    /// at a given layer. The Merkle tree is built by first transposing evaluations into a
    /// two-dimensional matrix where each row contains values needed to compute a single
    /// value of the next FRI layer, and then putting each row of the matrix into a single
    /// leaf of the Merkle tree. Thus, the number of elements grouped into a single leaf is
    /// equal to the `folding_factor` used for FRI layer construction.
    fn commit_fri_layer(
        &mut self,
        layer_root: <<Self as ProverChannel<E>>::Hasher as Hasher>::Digest,
    );

    /// Returns a random Î± drawn uniformly at random from the entire field.
    ///
    /// The prover uses this Î± to build the next FRI layer.
    ///
    /// While in the interactive version of the protocol the verifier send a random Î± to the
    /// prover, in the non-interactive version, the Î± is pseudo-randomly generated based on the
    /// values the prover previously wrote into the channel.
    fn draw_fri_alpha(&mut self) -> E;
}

// DEFAULT PROVER CHANNEL IMPLEMENTATION
// ================================================================================================

/// Provides a default implementation of the [ProverChannel] trait.
///
/// Though this implementation is intended primarily for testing purposes, it can be used in
/// production use cases as well.
pub struct DefaultProverChannel<E, H, R>
where
    E: FieldElement,
    H: ElementHasher<BaseField = E::BaseField>,
    R: RandomCoin<BaseField = E::BaseField, Hasher = H>,
{
    public_coin: R,
    commitments: Vec<H::Digest>,
    domain_size: usize,
    num_queries: usize,
    _field_element: PhantomData<E>,
}

impl<E, H, R> DefaultProverChannel<E, H, R>
where
    E: FieldElement,
    H: ElementHasher<BaseField = E::BaseField>,
    R: RandomCoin<BaseField = E::BaseField, Hasher = H>,
{
    /// Returns a new prover channel instantiated from the specified parameters.
    ///
    /// # Panics
    /// Panics if:
    /// * `domain_size` is smaller than 8 or is not a power of two.
    /// * `num_queries` is zero.
    pub fn new(domain_size: usize, num_queries: usize) -> Self {
        assert!(
            domain_size >= 8,
            "domain size must be at least 8, but was {domain_size}"
        );
        assert!(
            domain_size.is_power_of_two(),
            "domain size must be a power of two, but was {domain_size}"
        );
        assert!(
            num_queries > 0,
            "number of queries must be greater than zero"
        );
        DefaultProverChannel {
            public_coin: RandomCoin::new(&[]),
            commitments: Vec::new(),
            domain_size,
            num_queries,
            _field_element: PhantomData,
        }
    }

    /// Draws a set of positions at which the polynomial evaluations committed at the first FRI
    /// layer should be queried.
    ///
    /// The positions are pseudo-randomly generated based on the values the prover has written
    /// into this channel.
    ///
    /// # Panics
    /// Panics if the specified number of unique positions could not be drawn from the specified
    /// domain. Both number of queried positions and domain size are specified during
    /// construction of the channel.
    pub fn draw_query_positions(&mut self) -> Vec<usize> {
        self.public_coin
            .draw_integers(self.num_queries, self.domain_size)
            .expect("failed to draw query position")
    }

    /// Returns a list of FRI layer commitments written by the prover into this channel.
    pub fn layer_commitments(&self) -> &[H::Digest] {
        &self.commitments
    }
}

impl<E, H, R> ProverChannel<E> for DefaultProverChannel<E, H, R>
where
    E: FieldElement,
    H: ElementHasher<BaseField = E::BaseField>,
    R: RandomCoin<BaseField = E::BaseField, Hasher = H>,
{
    type Hasher = H;

    fn commit_fri_layer(&mut self, layer_root: H::Digest) {
        self.commitments.push(layer_root);
        self.public_coin.reseed(layer_root);
    }

    fn draw_fri_alpha(&mut self) -> E {
        self.public_coin.draw().expect("failed to draw FRI alpha")
    }
}


// Copyright (c) Facebook, Inc. and its affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use crate::{
    folding::{apply_drp, fold_positions},
    proof::{FriProof, FriProofLayer},
    utils::hash_values,
    FriOptions,
};
use core::marker::PhantomData;
use crypto::{ElementHasher, Hasher, MerkleTree};
use math::{fft, FieldElement, StarkField};
use utils::{collections::Vec, flatten_vector_elements, group_slice_elements, transpose_slice};

mod channel;
pub use channel::{DefaultProverChannel, ProverChannel};

#[cfg(test)]
mod tests;

// TYPES AND INTERFACES
// ================================================================================================

/// Implements the prover component of the FRI protocol.
///
/// Given evaluations of a function *f* over domain *D* (`evaluations`), a FRI prover generates
/// a proof that *f* is a polynomial of some bounded degree *d*, such that *d* < |*D*| / *blowup_factor*.
/// The proof is succinct: it exponentially smaller than `evaluations` and the verifier can verify it
/// exponentially faster than it would have taken them to read all `evaluations`.
///
/// The prover is parametrized with the following types:
///
/// * `B` specifies the base field of the STARK protocol.
/// * `E` specifies the field in which the FRI protocol is executed. This can be the same as the
///   base field `B`, but it can also be an extension of the base field in cases when the base
///   field is too small to provide desired security level for the FRI protocol.
/// * `C` specifies the type used to simulate prover-verifier interaction.
/// * `H` specifies the hash function used to build layer Merkle trees. The same hash function
///   must be used in the prover channel to generate pseudo random values.
///
/// Proof generation is performed in two phases: commit phase and query phase.
///
/// # Commit phase
/// During the commit phase, which is executed via [build_layers()](FriProver::build_layers())
/// function, the prover repeatedly applies a degree-respecting projection (DRP) to `evaluations`
/// (see [folding](crate::folding)). With every application of the DRP, the degree of the function
/// *f* (and size of the domain over which it is evaluated) is reduced by the `folding_factor`
/// until the remaining evaluations correspond to a polynomial, called remainder polynomial, with
/// a number of coefficients less than or equal to `remainder_max_degree_plus_1`.
///
/// At each layer of reduction, the prover commits to the current set of evaluations. This is done
/// by building a Merkle tree from the evaluations and sending the root of the tree to the verifier
/// (via [ProverChannel]). The Merkle tree is build in such a way that all evaluations needed to
/// compute a single value in the next FRI layer are grouped into the same leaf (the number of
/// evaluations needed to compute a single element in the next FRI layer is equal to the
/// `folding_factor`). This allows us to decommit all these values using a single Merkle
/// authentication path.
///
/// After committing to the set of evaluations at the current layer, the prover draws a random
/// field element Î± from the channel, and uses it to build the next FRI layer. In the interactive
/// version of the protocol, the verifier draws Î± uniformly at random from the entire field and
/// sends it to the prover. In the non-interactive version, Î± is pseudo-randomly generated based
/// on the values the prover has written into the channel up to that point.
///
/// The prover keeps all FRI layers (consisting of evaluations and corresponding Merkle trees) in
/// its internal state.
///
/// # Query phase
/// In the query phase, which is executed via [build_proof()](FriProver::build_proof()) function,
/// the prover receives a set of positions in the domain *D* from the verifier. The prover then
/// decommits evaluations corresponding to these positions across all FRI layers (except for the
/// remainder layer) and builds a [FriProof] from these evaluations. The remainder polynomial
/// is included in the proof in its entirety.
///
/// In the interactive version of the protocol, the verifier draws the position uniformly at
/// random from domain *D*. In the non-interactive version, the positions are pseudo-randomly
/// selected based on the values the prover has written into the channel up to that point.
///
/// Since the positions are drawn from domain *D*, they apply directly only to the first FRI
/// layer. To map these positions to the positions in all subsequent layers, the prover uses
/// [fold_positions] procedure.
///
/// After the proof is generated, the prover deletes all internally stored FRI layers.
///
/// Calling [build_layers()](FriProver::build_layers()) when the internal state is dirty, or
/// calling [build_proof()](FriProver::build_proof()) on a clean state will result in a panic.
pub struct FriProver<B, E, C, H>
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
    C: ProverChannel<E, Hasher = H>,
    H: ElementHasher<BaseField = B>,
{
    options: FriOptions,
    layers: Vec<FriLayer<B, E, H>>,
    remainder_poly: FriRemainder<E>,
    _channel: PhantomData<C>,
}

struct FriLayer<B: StarkField, E: FieldElement<BaseField = B>, H: Hasher> {
    tree: MerkleTree<H>,
    evaluations: Vec<E>,
    _base_field: PhantomData<B>,
}

struct FriRemainder<E: FieldElement>(Vec<E>);

// PROVER IMPLEMENTATION
// ================================================================================================

impl<B, E, C, H> FriProver<B, E, C, H>
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
    C: ProverChannel<E, Hasher = H>,
    H: ElementHasher<BaseField = B>,
{
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new FRI prover instantiated with the provided `options`.
    pub fn new(options: FriOptions) -> Self {
        FriProver {
            options,
            layers: Vec::new(),
            remainder_poly: FriRemainder(vec![]),
            _channel: PhantomData,
        }
    }

    // ACCESSORS
    // --------------------------------------------------------------------------------------------

    /// Returns folding factor for this prover.
    pub fn folding_factor(&self) -> usize {
        self.options.folding_factor()
    }

    /// Returns offset of the domain over which FRI protocol is executed by this prover.
    pub fn domain_offset(&self) -> B {
        self.options.domain_offset()
    }

    /// Returns number of FRI layers computed during the last execution of the
    /// [build_layers()](FriProver::build_layers()) method.
    pub fn num_layers(&self) -> usize {
        self.layers.len()
    }

    /// Clears a vector of internally stored layers.
    pub fn reset(&mut self) {
        self.layers.clear();
        self.remainder_poly.0.clear();
    }

    // COMMIT PHASE
    // --------------------------------------------------------------------------------------------
    /// Executes the commit phase of the FRI protocol.
    ///
    /// During this phase we repeatedly apply a degree-respecting projection (DRP) to
    /// `evaluations` which contain evaluations of some function *f* over domain *D*. With every
    /// application of the DRP the degree of the function (and size of the domain) is reduced by
    /// `folding_factor` until the remaining evaluations can be represented by a remainder polynomial
    /// with at most `remainder_max_degree_plus_1` number of coefficients.
    /// At each layer of reduction the current evaluations are committed to using a Merkle tree,
    /// and the root of this tree is written into the channel. After this the prover draws a random
    /// field element Î± from the channel, and uses it in the next application of the DRP.
    ///
    /// # Panics
    /// Panics if the prover state is dirty (the vector of layers is not empty).
    pub fn build_layers(&mut self, channel: &mut C, mut evaluations: Vec<E>) {
        assert!(
            self.layers.is_empty(),
            "a prior proof generation request has not been completed yet"
        );

        // reduce the degree by folding_factor at each iteration until the remaining polynomial
        // has small enough degree
        for _ in 0..self.options.num_fri_layers(evaluations.len()) {
            match self.folding_factor() {
                2 => self.build_layer::<2>(channel, &mut evaluations),
                4 => self.build_layer::<4>(channel, &mut evaluations),
                8 => self.build_layer::<8>(channel, &mut evaluations),
                16 => self.build_layer::<16>(channel, &mut evaluations),
                _ => unimplemented!("folding factor {} is not supported", self.folding_factor()),
            }
        }

        self.set_remainder(channel, &mut evaluations);
    }

    /// Builds a single FRI layer by first committing to the `evaluations`, then drawing a random
    /// alpha from the channel and use it to perform degree-respecting projection.
    fn build_layer<const N: usize>(&mut self, channel: &mut C, evaluations: &mut Vec<E>) {
        // commit to the evaluations at the current layer; we do this by first transposing the
        // evaluations into a matrix of N columns, and then building a Merkle tree from the
        // rows of this matrix; we do this so that we could de-commit to N values with a single
        // Merkle authentication path.
        let transposed_evaluations = transpose_slice(evaluations);
        let hashed_evaluations = hash_values::<H, E, N>(&transposed_evaluations);
        let evaluation_tree =
            MerkleTree::<H>::new(hashed_evaluations).expect("failed to construct FRI layer tree");
        channel.commit_fri_layer(*evaluation_tree.root());

        // draw a pseudo-random coefficient from the channel, and use it in degree-respecting
        // projection to reduce the degree of evaluations by N
        let alpha = channel.draw_fri_alpha();
        *evaluations = apply_drp(&transposed_evaluations, self.domain_offset(), alpha);
        self.layers.push(FriLayer {
            tree: evaluation_tree,
            evaluations: flatten_vector_elements(transposed_evaluations),
            _base_field: PhantomData,
        });
    }

    /// Creates remainder polynomial in coefficient form from a vector of `evaluations` over a domain.
    fn set_remainder(&mut self, channel: &mut C, evaluations: &mut [E]) {
        let inv_twiddles = fft::get_inv_twiddles(evaluations.len());
        fft::interpolate_poly_with_offset(evaluations, &inv_twiddles, self.options.domain_offset());
        let remainder_poly_size = evaluations.len() / self.options.blowup_factor();
        let remainder_poly = evaluations[..remainder_poly_size].to_vec();
        let commitment = <H as ElementHasher>::hash_elements(&remainder_poly);
        channel.commit_fri_layer(commitment);
        self.remainder_poly = FriRemainder(remainder_poly);
    }

    // QUERY PHASE
    // --------------------------------------------------------------------------------------------
    /// Executes query phase of FRI protocol.
    ///
    /// For each of the provided `positions`, corresponding evaluations from each of the layers
    /// (excluding the remainder layer) are recorded into the proof together with Merkle
    /// authentication paths from the root of layer commitment trees. For the remainder, we send
    /// the whole remainder polynomial resulting from interpolating the remainder layer.
    ///
    /// # Panics
    /// Panics is the prover state is clean (no FRI layers have been build yet).
    pub fn build_proof(&mut self, positions: &[usize]) -> FriProof {
        assert!(
            !self.remainder_poly.0.is_empty(),
            "FRI layers have not been built yet"
        );

        let mut layers = Vec::with_capacity(self.layers.len());

        if !self.layers.is_empty() {
            let mut positions = positions.to_vec();
            let mut domain_size = self.layers[0].evaluations.len();
            let folding_factor = self.options.folding_factor();

            // for all FRI layers, except the last one, record tree root, determine a set of query
            // positions, and query the layer at these positions.
            for i in 0..self.layers.len() {
                positions = fold_positions(&positions, domain_size, folding_factor);

                // sort of a static dispatch for folding_factor parameter
                let proof_layer = match folding_factor {
                    2 => query_layer::<B, E, H, 2>(&self.layers[i], &positions),
                    4 => query_layer::<B, E, H, 4>(&self.layers[i], &positions),
                    8 => query_layer::<B, E, H, 8>(&self.layers[i], &positions),
                    16 => query_layer::<B, E, H, 16>(&self.layers[i], &positions),
                    _ => unimplemented!("folding factor {} is not supported", folding_factor),
                };

                layers.push(proof_layer);
                domain_size /= folding_factor;
            }
        }

        // use the remaining polynomial values directly as proof
        let remainder = self.remainder_poly.0.clone();

        // clear layers so that another proof can be generated
        self.reset();

        FriProof::new(layers, remainder, 1)
    }
}

// HELPER FUNCTIONS
// ================================================================================================

/// Builds a single proof layer by querying the evaluations of the passed in FRI layer at the
/// specified positions.
fn query_layer<B: StarkField, E: FieldElement<BaseField = B>, H: Hasher, const N: usize>(
    layer: &FriLayer<B, E, H>,
    positions: &[usize],
) -> FriProofLayer {
    // build Merkle authentication paths for all query positions
    let proof = layer
        .tree
        .prove_batch(positions)
        .expect("failed to generate a Merkle proof for FRI layer queries");

    // build a list of polynomial evaluations at each position; since evaluations in FRI layers
    // are stored in transposed form, a position refers to N evaluations which are committed
    // in a single leaf
    let evaluations: &[[E; N]] = group_slice_elements(&layer.evaluations);
    let mut queried_values: Vec<[E; N]> = Vec::with_capacity(positions.len());
    for &position in positions.iter() {
        queried_values.push(evaluations[position]);
    }

    FriProofLayer::new(queried_values, proof)
}


// Copyright (c) Facebook, Inc. and its affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use super::{DefaultProverChannel, FriProver};
use crate::{
    verifier::{DefaultVerifierChannel, FriVerifier},
    FriOptions, FriProof, VerifierError,
};
use crypto::{hashers::Blake3_256, DefaultRandomCoin, Hasher, RandomCoin};
use math::{fft, fields::f128::BaseElement, FieldElement};
use utils::{collections::Vec, Deserializable, Serializable, SliceReader};

type Blake3 = Blake3_256<BaseElement>;

// PROVE/VERIFY TEST
// ================================================================================================

#[test]
fn fri_folding_2() {
    let trace_length_e = 12;
    let lde_blowup_e = 3;
    let folding_factor_e = 1;
    let max_remainder_degree = 7;
    fri_prove_verify(
        trace_length_e,
        lde_blowup_e,
        folding_factor_e,
        max_remainder_degree,
    )
}

#[test]
fn fri_folding_4() {
    let trace_length_e = 12;
    let lde_blowup_e = 3;
    let folding_factor_e = 2;
    let max_remainder_degree = 255;
    fri_prove_verify(
        trace_length_e,
        lde_blowup_e,
        folding_factor_e,
        max_remainder_degree,
    )
}

// TEST UTILS
// ================================================================================================

pub fn build_prover_channel(
    trace_length: usize,
    options: &FriOptions,
) -> DefaultProverChannel<BaseElement, Blake3, DefaultRandomCoin<Blake3>> {
    DefaultProverChannel::new(trace_length * options.blowup_factor(), 32)
}

pub fn build_evaluations(trace_length: usize, lde_blowup: usize) -> Vec<BaseElement> {
    let mut p = (0..trace_length as u128)
        .map(BaseElement::new)
        .collect::<Vec<_>>();
    let domain_size = trace_length * lde_blowup;
    p.resize(domain_size, BaseElement::ZERO);

    let twiddles = fft::get_twiddles::<BaseElement>(domain_size);

    fft::evaluate_poly(&mut p, &twiddles);
    p
}

pub fn verify_proof(
    proof: FriProof,
    commitments: Vec<<Blake3 as Hasher>::Digest>,
    evaluations: &[BaseElement],
    max_degree: usize,
    domain_size: usize,
    positions: &[usize],
    options: &FriOptions,
) -> Result<(), VerifierError> {
    // test proof serialization / deserialization
    let mut proof_bytes = Vec::new();
    proof.write_into(&mut proof_bytes);

    let mut reader = SliceReader::new(&proof_bytes);
    let proof = FriProof::read_from(&mut reader).unwrap();

    // verify the proof
    let mut channel = DefaultVerifierChannel::<BaseElement, Blake3>::new(
        proof,
        commitments,
        domain_size,
        options.folding_factor(),
    )
    .unwrap();
    let mut coin = DefaultRandomCoin::<Blake3>::new(&[]);
    let verifier = FriVerifier::new(&mut channel, &mut coin, options.clone(), max_degree)?;
    let queried_evaluations = positions
        .iter()
        .map(|&p| evaluations[p])
        .collect::<Vec<_>>();
    verifier.verify(&mut channel, &queried_evaluations, positions)
}

fn fri_prove_verify(
    trace_length_e: usize,
    lde_blowup_e: usize,
    folding_factor_e: usize,
    max_remainder_degree: usize,
) {
    let trace_length = 1 << trace_length_e;
    let lde_blowup = 1 << lde_blowup_e;
    let folding_factor = 1 << folding_factor_e;

    let options = FriOptions::new(lde_blowup, folding_factor, max_remainder_degree);
    let mut channel = build_prover_channel(trace_length, &options);
    let evaluations = build_evaluations(trace_length, lde_blowup);

    // instantiate the prover and generate the proof
    let mut prover = FriProver::new(options.clone());
    prover.build_layers(&mut channel, evaluations.clone());
    let positions = channel.draw_query_positions();
    let proof = prover.build_proof(&positions);

    // make sure the proof can be verified
    let commitments = channel.layer_commitments().to_vec();
    let max_degree = trace_length - 1;
    let result = verify_proof(
        proof.clone(),
        commitments.clone(),
        &evaluations,
        max_degree,
        trace_length * lde_blowup,
        &positions,
        &options,
    );
    assert!(result.is_ok(), "{:}", result.err().unwrap());

    // make sure proof fails for invalid degree
    let result = verify_proof(
        proof,
        commitments,
        &evaluations,
        max_degree - 8,
        trace_length * lde_blowup,
        &positions,
        &options,
    );
    assert!(result.is_err());
}


// Copyright (c) Facebook, Inc. and its affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use crate::{FriProof, VerifierError};
use crypto::{BatchMerkleProof, ElementHasher, Hasher, MerkleTree};
use math::FieldElement;
use utils::{collections::Vec, group_vector_elements, DeserializationError};

// VERIFIER CHANNEL TRAIT
// ================================================================================================

/// Defines an interface for a channel over which a verifier communicates with a prover.
///
/// This trait abstracts away implementation specifics of the [FriProof] struct. Thus, instead of
/// dealing with FRI proofs directly, the verifier can read the data as if it was sent by the
/// prover via an interactive channel.
///
/// Note: that reading removes the data from the channel. Thus, reading duplicated values from
/// the channel should not be possible.
pub trait VerifierChannel<E: FieldElement> {
    /// Hash function used by the prover to commit to polynomial evaluations.
    type Hasher: ElementHasher<BaseField = E::BaseField>;

    // REQUIRED METHODS
    // --------------------------------------------------------------------------------------------

    /// Returns the number of partitions used during proof generation.
    fn read_fri_num_partitions(&self) -> usize;

    /// Reads and removes from the channel all FRI layer commitments sent by the prover.
    ///
    /// In the interactive version of the protocol, the prover sends layer commitments to the
    /// verifier one-by-one, and the verifier responds with a value Î± drawn uniformly at random
    /// from the entire field after each layer commitment is received. In the non-interactive
    /// version, the verifier can read all layer commitments at once, and then generate Î± values
    /// locally.
    fn read_fri_layer_commitments(
        &mut self,
    ) -> Vec<<<Self as VerifierChannel<E>>::Hasher as Hasher>::Digest>;

    /// Reads and removes from the channel evaluations of the polynomial at the queried positions
    /// for the next FRI layer.
    ///
    /// In the interactive version of the protocol, these evaluations are sent from the prover to
    /// the verifier during the query phase of the FRI protocol.
    ///
    /// It is expected that layer queries and layer proofs at the same FRI layer are consistent.
    /// That is, query values hash into the leaf nodes of corresponding Merkle authentication
    /// paths.
    fn take_next_fri_layer_queries(&mut self) -> Vec<E>;

    /// Reads and removes from the channel Merkle authentication paths for queried evaluations for
    /// the next FRI layer.
    ///
    /// In the interactive version of the protocol, these authentication paths are sent from the
    /// prover to the verifier during the query phase of the FRI protocol.
    ///
    /// It is expected that layer proofs and layer queries at the same FRI layer are consistent.
    /// That is, query values hash into the leaf nodes of corresponding Merkle authentication
    /// paths.
    fn take_next_fri_layer_proof(&mut self) -> BatchMerkleProof<Self::Hasher>;

    /// Reads and removes the remainder polynomial from the channel.
    fn take_fri_remainder(&mut self) -> Vec<E>;

    // PROVIDED METHODS
    // --------------------------------------------------------------------------------------------

    /// Returns FRI query values at the specified positions from the current FRI layer and advances
    /// layer pointer by one.
    ///
    /// This also checks if the values are valid against the provided FRI layer commitment.
    ///
    /// # Errors
    /// Returns an error if query values did not match layer commitment.
    fn read_layer_queries<const N: usize>(
        &mut self,
        positions: &[usize],
        commitment: &<<Self as VerifierChannel<E>>::Hasher as Hasher>::Digest,
    ) -> Result<Vec<[E; N]>, VerifierError> {
        let layer_proof = self.take_next_fri_layer_proof();
        MerkleTree::<Self::Hasher>::verify_batch(commitment, positions, &layer_proof)
            .map_err(|_| VerifierError::LayerCommitmentMismatch)?;

        // TODO: make sure layer queries hash into leaves of layer proof

        let layer_queries = self.take_next_fri_layer_queries();
        Ok(group_vector_elements(layer_queries))
    }

    /// Returns FRI remainder polynomial read from this channel.
    fn read_remainder(&mut self) -> Result<Vec<E>, VerifierError> {
        let remainder = self.take_fri_remainder();

        Ok(remainder)
    }
}

// DEFAULT VERIFIER CHANNEL IMPLEMENTATION
// ================================================================================================

/// Provides a default implementation of the [VerifierChannel] trait.
///
/// Default verifier channel can be instantiated directly from a [FriProof] struct.
///
/// Though this implementation is primarily intended for testing purposes, it can be used in
/// production use cases as well.
pub struct DefaultVerifierChannel<E: FieldElement, H: ElementHasher<BaseField = E::BaseField>> {
    layer_commitments: Vec<H::Digest>,
    layer_proofs: Vec<BatchMerkleProof<H>>,
    layer_queries: Vec<Vec<E>>,
    remainder: Vec<E>,
    num_partitions: usize,
}

impl<E, H> DefaultVerifierChannel<E, H>
where
    E: FieldElement,
    H: ElementHasher<BaseField = E::BaseField>,
{
    /// Builds a new verifier channel from the specified [FriProof].
    ///
    /// # Errors
    /// Returns an error if the specified `proof` could not be parsed correctly.
    pub fn new(
        proof: FriProof,
        layer_commitments: Vec<H::Digest>,
        domain_size: usize,
        folding_factor: usize,
    ) -> Result<Self, DeserializationError> {
        let num_partitions = proof.num_partitions();

        let remainder = proof.parse_remainder()?;
        let (layer_queries, layer_proofs) =
            proof.parse_layers::<H, E>(domain_size, folding_factor)?;

        Ok(DefaultVerifierChannel {
            layer_commitments,
            layer_proofs,
            layer_queries,
            remainder,
            num_partitions,
        })
    }
}

impl<E, H> VerifierChannel<E> for DefaultVerifierChannel<E, H>
where
    E: FieldElement,
    H: ElementHasher<BaseField = E::BaseField>,
{
    type Hasher = H;

    fn read_fri_num_partitions(&self) -> usize {
        self.num_partitions
    }

    fn read_fri_layer_commitments(&mut self) -> Vec<H::Digest> {
        self.layer_commitments.drain(..).collect()
    }

    fn take_next_fri_layer_proof(&mut self) -> BatchMerkleProof<H> {
        self.layer_proofs.remove(0)
    }

    fn take_next_fri_layer_queries(&mut self) -> Vec<E> {
        self.layer_queries.remove(0)
    }

    fn take_fri_remainder(&mut self) -> Vec<E> {
        self.remainder.clone()
    }
}


// Copyright (c) Facebook, Inc. and its affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

//! Contains an implementation of FRI verifier and associated components.

use crate::{folding::fold_positions, utils::map_positions_to_indexes, FriOptions, VerifierError};
use core::{convert::TryInto, marker::PhantomData, mem};
use crypto::{ElementHasher, RandomCoin};
use math::{polynom, FieldElement, StarkField};
use utils::collections::Vec;

mod channel;
pub use channel::{DefaultVerifierChannel, VerifierChannel};

// FRI VERIFIER
// ================================================================================================
/// Implements the verifier component of the FRI protocol.
///
/// Given a small number of evaluations of some function *f* over domain *D* and a FRI proof, a
/// FRI verifier determines whether *f* is a polynomial of some bounded degree *d*, such that *d*
/// < |*D*| / 2.
///
/// The verifier is parametrized by the following types:
///
/// * `B` specifies the base field of the STARK protocol.
/// * `E` specifies the field in which the FRI protocol is executed. This can be the same as the
///   base field `B`, but it can also be an extension of the base field in cases when the base
///   field is too small to provide desired security level for the FRI protocol.
/// * `C` specifies the type used to simulate prover-verifier interaction. This type is used
///   as an abstraction for a [FriProof](crate::FriProof). Meaning, the verifier does not consume
///   a FRI proof directly, but reads it via [VerifierChannel] interface.
/// * `H` specifies the Hash function used by the prover to commit to polynomial evaluations.
///
/// Proof verification is performed in two phases: commit phase and query phase.
///
/// # Commit phase
/// During the commit phase, which is executed when the verifier is instantiated via
/// [new()](FriVerifier::new()) function, the verifier receives a list of FRI layer commitments
/// from the prover (via [VerifierChannel]). After each received commitment, the verifier
/// draws a random value Î± from the entire field, and sends it to the prover. In the
/// non-interactive version of the protocol, Î± values are derived pseudo-randomly from FRI
/// layer commitments.
///
/// # Query phase
/// During the query phase, which is executed via [verify()](FriVerifier::verify()) function,
/// the verifier sends a set of positions in the domain *D* to the prover, and the prover responds
/// with polynomial evaluations at these positions (together with corresponding Merkle paths)
/// across all FRI layers. The verifier then checks that:
/// * The Merkle paths are valid against the layer commitments the verifier received during
///   the commit phase.
/// * The evaluations are consistent across FRI layers (i.e., the degree-respecting projection
///   was applied correctly).
/// * The degree of the polynomial implied by evaluations at the last FRI layer (the remainder)
///   is smaller than the degree resulting from reducing degree *d* by `folding_factor` at each
///   FRI layer.
pub struct FriVerifier<E, C, H, R>
where
    E: FieldElement,
    C: VerifierChannel<E, Hasher = H>,
    H: ElementHasher<BaseField = E::BaseField>,
    R: RandomCoin<BaseField = E::BaseField, Hasher = H>,
{
    max_poly_degree: usize,
    domain_size: usize,
    domain_generator: E::BaseField,
    layer_commitments: Vec<H::Digest>,
    layer_alphas: Vec<E>,
    options: FriOptions,
    num_partitions: usize,
    _channel: PhantomData<C>,
    _public_coin: PhantomData<R>,
}

impl<E, C, H, R> FriVerifier<E, C, H, R>
where
    E: FieldElement,
    C: VerifierChannel<E, Hasher = H>,
    H: ElementHasher<BaseField = E::BaseField>,
    R: RandomCoin<BaseField = E::BaseField, Hasher = H>,
{
    /// Returns a new instance of FRI verifier created from the specified parameters.
    ///
    /// The `max_poly_degree` parameter specifies the highest polynomial degree accepted by the
    /// returned verifier. In combination with `blowup_factor` from the `options` parameter,
    /// `max_poly_degree` also defines the domain over which the tested polynomial is evaluated.
    ///
    /// Creating a FRI verifier executes the commit phase of the FRI protocol from the verifier's
    /// perspective. Specifically, the verifier reads FRI layer commitments from the `channel`,
    /// and for each commitment, updates the `public_coin` with this commitment and then draws
    /// a random value Î± from the coin.
    ///
    /// The verifier stores layer commitments and corresponding Î± values in its internal state,
    /// and, thus, an instance of FRI verifier can be used to verify only a single proof.
    ///
    /// # Errors
    /// Returns an error if:
    /// * `max_poly_degree` is inconsistent with the number of FRI layers read from the channel
    ///   and `folding_factor` specified in the `options` parameter.
    /// * An error was encountered while drawing a random Î± value from the coin.
    pub fn new(
        channel: &mut C,
        public_coin: &mut R,
        options: FriOptions,
        max_poly_degree: usize,
    ) -> Result<Self, VerifierError> {
        // infer evaluation domain info
        let domain_size = max_poly_degree.next_power_of_two() * options.blowup_factor();
        let domain_generator = E::BaseField::get_root_of_unity(domain_size.ilog2());

        let num_partitions = channel.read_fri_num_partitions();

        // read layer commitments from the channel and use them to build a list of alphas
        let layer_commitments = channel.read_fri_layer_commitments();
        let mut layer_alphas = Vec::with_capacity(layer_commitments.len());
        let mut max_degree_plus_1 = max_poly_degree + 1;
        for (depth, commitment) in layer_commitments.iter().enumerate() {
            public_coin.reseed(*commitment);
            let alpha = public_coin.draw().map_err(VerifierError::RandomCoinError)?;
            layer_alphas.push(alpha);

            // make sure the degree can be reduced by the folding factor at all layers
            // but the remainder layer
            if depth != layer_commitments.len() - 1
                && max_degree_plus_1 % options.folding_factor() != 0
            {
                return Err(VerifierError::DegreeTruncation(
                    max_degree_plus_1 - 1,
                    options.folding_factor(),
                    depth,
                ));
            }
            max_degree_plus_1 /= options.folding_factor();
        }

        Ok(FriVerifier {
            max_poly_degree,
            domain_size,
            domain_generator,
            layer_commitments,
            layer_alphas,
            options,
            num_partitions,
            _channel: PhantomData,
            _public_coin: PhantomData,
        })
    }

    // PUBLIC ACCESSORS
    // --------------------------------------------------------------------------------------------

    /// Returns maximum degree of a polynomial accepted by this verifier.
    pub fn max_poly_degree(&self) -> usize {
        self.max_poly_degree
    }

    /// Returns size of the domain over which a polynomial commitment checked by this verifier
    /// has been evaluated.
    ///
    /// The domain size can be computed by rounding `max_poly_degree` to the next power of two
    /// and multiplying the result by the `blowup_factor` from the protocol options.
    pub fn domain_size(&self) -> usize {
        self.domain_size
    }

    /// Returns number of partitions used during FRI proof generation.
    ///
    /// For non-distributed proof generation, number of partitions is usually set to 1.
    pub fn num_partitions(&self) -> usize {
        self.num_partitions
    }

    /// Returns protocol configuration options for this verifier.
    pub fn options(&self) -> &FriOptions {
        &self.options
    }

    // VERIFICATION PROCEDURE
    // --------------------------------------------------------------------------------------------
    /// Executes the query phase of the FRI protocol.
    ///
    /// Returns `Ok(())` if values in the `evaluations` slice represent evaluations of a polynomial
    /// with degree <= `max_poly_degree` at x coordinates specified by the `positions` slice.
    ///
    /// Thus, `positions` parameter represents the positions in the evaluation domain at which the
    /// verifier queries the prover at the first FRI layer. Similarly, the `evaluations` parameter
    /// specifies the evaluations of the polynomial at the first FRI layer returned by the prover
    /// for these positions.
    ///
    /// Evaluations of layer polynomials for all subsequent FRI layers the verifier reads from the
    /// specified `channel`.
    ///
    /// # Errors
    /// Returns an error if:
    /// * The length of `evaluations` is not equal to the length of `positions`.
    /// * An unsupported folding factor was specified by the `options` for this verifier.
    /// * Decommitments to polynomial evaluations don't match the commitment value at any of the
    ///   FRI layers.
    /// * The verifier detects an error in how the degree-respecting projection was applied
    ///   at any of the FRI layers.
    /// * The degree of the remainder at the last FRI layer is greater than the degree implied by
    ///   `max_poly_degree` reduced by the folding factor at each FRI layer.
    pub fn verify(
        &self,
        channel: &mut C,
        evaluations: &[E],
        positions: &[usize],
    ) -> Result<(), VerifierError> {
        if evaluations.len() != positions.len() {
            return Err(VerifierError::NumPositionEvaluationMismatch(
                positions.len(),
                evaluations.len(),
            ));
        }

        // static dispatch for folding factor parameter
        let folding_factor = self.options.folding_factor();
        match folding_factor {
            2 => self.verify_generic::<2>(channel, evaluations, positions),
            4 => self.verify_generic::<4>(channel, evaluations, positions),
            8 => self.verify_generic::<8>(channel, evaluations, positions),
            16 => self.verify_generic::<16>(channel, evaluations, positions),
            _ => Err(VerifierError::UnsupportedFoldingFactor(folding_factor)),
        }
    }

    /// This is the actual implementation of the verification procedure described above, but it
    /// also takes folding factor as a generic parameter N.
    fn verify_generic<const N: usize>(
        &self,
        channel: &mut C,
        evaluations: &[E],
        positions: &[usize],
    ) -> Result<(), VerifierError> {
        // pre-compute roots of unity used in computing x coordinates in the folded domain
        let folding_roots = (0..N)
            .map(|i| {
                self.domain_generator
                    .exp_vartime(((self.domain_size / N * i) as u64).into())
            })
            .collect::<Vec<_>>();

        // 1 ----- verify the recursive components of the FRI proof -----------------------------------
        let mut domain_generator = self.domain_generator;
        let mut domain_size = self.domain_size;
        let mut max_degree_plus_1 = self.max_poly_degree + 1;
        let mut positions = positions.to_vec();
        let mut evaluations = evaluations.to_vec();

        for depth in 0..self.options.num_fri_layers(self.domain_size) {
            // determine which evaluations were queried in the folded layer
            let mut folded_positions =
                fold_positions(&positions, domain_size, self.options.folding_factor());
            // determine where these evaluations are in the commitment Merkle tree
            let position_indexes = map_positions_to_indexes(
                &folded_positions,
                domain_size,
                self.options.folding_factor(),
                self.num_partitions,
            );
            // read query values from the specified indexes in the Merkle tree
            let layer_commitment = self.layer_commitments[depth];
            // TODO: add layer depth to the potential error message
            let layer_values = channel.read_layer_queries(&position_indexes, &layer_commitment)?;
            let query_values =
                get_query_values::<E, N>(&layer_values, &positions, &folded_positions, domain_size);
            if evaluations != query_values {
                return Err(VerifierError::InvalidLayerFolding(depth));
            }

            // build a set of x coordinates for each row polynomial
            #[rustfmt::skip]
            let xs = folded_positions.iter().map(|&i| {
                let xe = domain_generator.exp_vartime((i as u64).into()) * self.options.domain_offset();
                folding_roots.iter()
                    .map(|&r| E::from(xe * r))
                    .collect::<Vec<_>>().try_into().unwrap()
            })
            .collect::<Vec<_>>();

            // interpolate x and y values into row polynomials
            let row_polys = polynom::interpolate_batch(&xs, &layer_values);

            // calculate the pseudo-random value used for linear combination in layer folding
            let alpha = self.layer_alphas[depth];

            // check that when the polynomials are evaluated at alpha, the result is equal to
            // the corresponding column value
            evaluations = row_polys.iter().map(|p| polynom::eval(p, alpha)).collect();

            // make sure next degree reduction does not result in degree truncation
            if max_degree_plus_1 % N != 0 {
                return Err(VerifierError::DegreeTruncation(
                    max_degree_plus_1 - 1,
                    N,
                    depth,
                ));
            }

            // update variables for the next iteration of the loop
            domain_generator = domain_generator.exp_vartime((N as u32).into());
            max_degree_plus_1 /= N;
            domain_size /= N;
            mem::swap(&mut positions, &mut folded_positions);
        }

        // 2 ----- verify the remainder polynomial of the FRI proof -------------------------------

        // read the remainder polynomial from the channel and make sure it agrees with the evaluations
        // from the previous layer.
        let remainder_poly = channel.read_remainder()?;
        if remainder_poly.len() > max_degree_plus_1 {
            return Err(VerifierError::RemainderDegreeMismatch(
                max_degree_plus_1 - 1,
            ));
        }
        let offset: E::BaseField = self.options().domain_offset();

        for (&position, evaluation) in positions.iter().zip(evaluations) {
            let comp_eval = eval_horner::<E>(
                &remainder_poly,
                offset * domain_generator.exp_vartime((position as u64).into()),
            );
            if comp_eval != evaluation {
                return Err(VerifierError::InvalidRemainderFolding);
            }
        }

        Ok(())
    }
}

// HELPER FUNCTIONS
// ================================================================================================
fn get_query_values<E: FieldElement, const N: usize>(
    values: &[[E; N]],
    positions: &[usize],
    folded_positions: &[usize],
    domain_size: usize,
) -> Vec<E> {
    let row_length = domain_size / N;

    let mut result = Vec::new();
    for position in positions {
        let idx = folded_positions
            .iter()
            .position(|&v| v == position % row_length)
            .unwrap();
        let value = values[idx][position / row_length];
        result.push(value);
    }

    result
}

// Evaluates a polynomial with coefficients in an extension field at a point in the base field.
pub fn eval_horner<E>(p: &[E], x: E::BaseField) -> E
where
    E: FieldElement,
{
    p.iter()
        .rev()
        .fold(E::ZERO, |acc, &coeff| acc * E::from(x) + coeff)
}


// Copyright (c) Facebook, Inc. and its affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
use rand_utils::rand_vector;
use std::time::Duration;
use winter_math::{
    fft,
    fields::{f128, f62, f64, CubeExtension, QuadExtension},
    FieldElement, StarkField,
};

const SIZES: [usize; 3] = [262_144, 524_288, 1_048_576];

fn fft_evaluate_poly<B, E>(c: &mut Criterion, field_name: &str)
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    let mut group = c.benchmark_group(format!("{field_name}/fft_evaluate_poly"));
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(10));

    let blowup_factor = 8;

    for &size in SIZES.iter() {
        let p: Vec<E> = rand_vector(size / blowup_factor);
        let twiddles: Vec<B> = fft::get_twiddles(size);
        group.bench_function(BenchmarkId::new("simple", size), |bench| {
            bench.iter_with_large_drop(|| {
                let mut result = vec![E::ZERO; size];
                result[..p.len()].copy_from_slice(&p);
                fft::evaluate_poly(&mut result, &twiddles);
                result
            });
        });
    }

    for &size in SIZES.iter() {
        let p: Vec<E> = rand_vector(size / blowup_factor);
        let twiddles: Vec<B> = fft::get_twiddles(size / blowup_factor);
        group.bench_function(BenchmarkId::new("with_offset", size), |bench| {
            bench.iter_with_large_drop(|| {
                fft::evaluate_poly_with_offset(&p, &twiddles, B::GENERATOR, blowup_factor)
            });
        });
    }

    group.finish();
}

fn fft_interpolate_poly<B, E>(c: &mut Criterion, field_name: &str)
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    let mut group = c.benchmark_group(format!("{field_name}/fft_interpolate_poly"));
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(10));

    for &size in SIZES.iter() {
        let p: Vec<E> = rand_vector(size);
        let inv_twiddles: Vec<B> = fft::get_inv_twiddles(size);
        group.bench_function(BenchmarkId::new("simple", size), |bench| {
            bench.iter_batched_ref(
                || p.clone(),
                |p| fft::interpolate_poly(p, &inv_twiddles),
                BatchSize::LargeInput,
            );
        });
    }

    for &size in SIZES.iter() {
        let p: Vec<E> = rand_vector(size);
        let inv_twiddles: Vec<B> = fft::get_inv_twiddles(size);
        group.bench_function(BenchmarkId::new("with_offset", size), |bench| {
            bench.iter_batched_ref(
                || p.clone(),
                |p| fft::interpolate_poly_with_offset(p, &inv_twiddles, B::GENERATOR),
                BatchSize::LargeInput,
            );
        });
    }

    group.finish();
}

fn get_twiddles(c: &mut Criterion) {
    let mut group = c.benchmark_group("fft_get_twiddles");
    group.sample_size(10);
    for &size in SIZES.iter() {
        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |bench, &size| {
            bench.iter(|| fft::get_twiddles::<f128::BaseElement>(size));
        });
    }
    group.finish();
}

fn bench_fft(c: &mut Criterion) {
    fft_evaluate_poly::<f62::BaseElement, f62::BaseElement>(c, "f62");
    fft_evaluate_poly::<f64::BaseElement, f64::BaseElement>(c, "f64");
    fft_evaluate_poly::<f128::BaseElement, f128::BaseElement>(c, "f128");

    fft_evaluate_poly::<f62::BaseElement, QuadExtension<f62::BaseElement>>(c, "f62_quad");
    fft_evaluate_poly::<f64::BaseElement, QuadExtension<f64::BaseElement>>(c, "f64_quad");
    fft_evaluate_poly::<f128::BaseElement, QuadExtension<f128::BaseElement>>(c, "f128_quad");

    fft_evaluate_poly::<f64::BaseElement, CubeExtension<f64::BaseElement>>(c, "f64_cube");

    fft_interpolate_poly::<f62::BaseElement, f62::BaseElement>(c, "f62");
    fft_interpolate_poly::<f64::BaseElement, f64::BaseElement>(c, "f64");
    fft_interpolate_poly::<f128::BaseElement, f128::BaseElement>(c, "f128");
}

criterion_group!(fft_group, bench_fft, get_twiddles);
criterion_main!(fft_group);


// Copyright (c) Facebook, Inc. and its affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use criterion::{
    black_box, criterion_group, criterion_main,
    measurement::{Measurement, WallTime},
    BatchSize, BenchmarkGroup, BenchmarkId, Criterion,
};
use rand_utils::{rand_array, rand_value, rand_vector};
use std::time::Duration;
use winter_math::{
    batch_inversion,
    fields::{f128, f62, f64},
    fields::{CubeExtension, QuadExtension},
    ExtensibleField, FieldElement, StarkField,
};

const SIZES: [usize; 3] = [262_144, 524_288, 1_048_576];

// BATCH INVERSION
// ================================================================================================

pub fn batch_inv(c: &mut Criterion) {
    let mut group = c.benchmark_group("batch_inv");
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(10));

    for &size in SIZES.iter() {
        let values = rand_vector::<f128::BaseElement>(size);

        group.bench_function(BenchmarkId::new("no_coeff", size), |bench| {
            bench.iter_with_large_drop(|| batch_inversion(&values));
        });
    }

    group.finish();
}

// SEQUENTIAL OPS
// ================================================================================================
pub fn field_ops<B>(c: &mut Criterion, field_name: &str)
where
    B: StarkField + ExtensibleField<2> + ExtensibleField<3>,
{
    let mut group = c.benchmark_group(format!("field/{field_name}"));

    // --- base field -----------------------------------------------------------------------------

    group.bench_function("add", |bench| {
        let x = rand_value::<B>();
        let y = rand_value::<B>();
        bench.iter(|| black_box(x) + black_box(y))
    });

    group.bench_function("double", |bench| {
        let x = rand_value::<B>();
        bench.iter(|| black_box(x).double())
    });

    group.bench_function("sub", |bench| {
        let x = rand_value::<B>();
        let y = rand_value::<B>();
        bench.iter(|| black_box(x) - black_box(y))
    });

    group.bench_function("mul", |bench| {
        let x = rand_value::<B>();
        let y = rand_value::<B>();
        bench.iter(|| black_box(x) * black_box(y))
    });

    group.bench_function("exp", |bench| {
        let x = rand_value::<B>();
        let y = rand_value::<B>().as_int();
        bench.iter(|| x.exp(y))
    });

    group.bench_function("inv", |bench| {
        let x = rand_value::<B>();
        bench.iter(|| x.inv())
    });

    batch_ops::<B, WallTime>(&mut group, "base");
    array_ops::<B, WallTime>(&mut group, "base");

    // --- quadratic extension --------------------------------------------------------------------

    if QuadExtension::<B>::is_supported() {
        group.bench_function("quad/add", |bench| {
            let x = rand_value::<QuadExtension<B>>();
            let y = rand_value::<QuadExtension<B>>();
            bench.iter(|| black_box(x) + black_box(y))
        });

        group.bench_function("quad/double", |bench| {
            let x = rand_value::<QuadExtension<B>>();
            bench.iter(|| black_box(x).double())
        });

        group.bench_function("quad/sub", |bench| {
            let x = rand_value::<QuadExtension<B>>();
            let y = rand_value::<QuadExtension<B>>();
            bench.iter(|| black_box(x) - black_box(y))
        });

        group.bench_function("quad/mul", |bench| {
            let x = rand_value::<QuadExtension<B>>();
            let y = rand_value::<QuadExtension<B>>();
            bench.iter(|| black_box(x) * black_box(y))
        });

        batch_ops::<QuadExtension<B>, WallTime>(&mut group, "quad");
        array_ops::<QuadExtension<B>, WallTime>(&mut group, "quad");
    }

    // --- cubic extension ------------------------------------------------------------------------

    if CubeExtension::<B>::is_supported() {
        group.bench_function("cube/add", |bench| {
            let x = rand_value::<CubeExtension<B>>();
            let y = rand_value::<CubeExtension<B>>();
            bench.iter(|| black_box(x) + black_box(y))
        });

        group.bench_function("cube/double", |bench| {
            let x = rand_value::<CubeExtension<B>>();
            bench.iter(|| black_box(x).double())
        });

        group.bench_function("cube/sub", |bench| {
            let x = rand_value::<CubeExtension<B>>();
            let y = rand_value::<CubeExtension<B>>();
            bench.iter(|| black_box(x) - black_box(y))
        });

        group.bench_function("cube/mul", |bench| {
            let x = rand_value::<CubeExtension<B>>();
            let y = rand_value::<CubeExtension<B>>();
            bench.iter(|| black_box(x) * black_box(y))
        });
    }
}

// ARRAY OPS
// ================================================================================================
pub fn array_ops<E: FieldElement, M: Measurement>(group: &mut BenchmarkGroup<M>, extension: &str) {
    group.bench_function(format!("{extension}/array/add"), |b| {
        b.iter_batched(
            || (rand_array::<E, 100>(), rand_array::<E, 100>()),
            |(mut x, y)| {
                for (x, y) in x.iter_mut().zip(y) {
                    *x += y;
                }
                x
            },
            BatchSize::SmallInput,
        )
    });

    group.bench_function(format!("{extension}/array/sub"), |b| {
        b.iter_batched(
            || (rand_array::<E, 100>(), rand_array::<E, 100>()),
            |(mut x, y)| {
                for (x, y) in x.iter_mut().zip(y) {
                    *x -= y;
                }
                x
            },
            BatchSize::SmallInput,
        )
    });

    group.bench_function(format!("{extension}/array/mul"), |b| {
        b.iter_batched(
            || (rand_array::<E, 100>(), rand_array::<E, 100>()),
            |(mut x, y)| {
                for (x, y) in x.iter_mut().zip(y) {
                    *x *= y;
                }
                x
            },
            BatchSize::SmallInput,
        )
    });
}

// BATCH OPS
// ================================================================================================
pub fn batch_ops<E: FieldElement, M: Measurement>(group: &mut BenchmarkGroup<M>, extension: &str) {
    group.bench_function(format!("{extension}/batch/add"), |b| {
        b.iter_batched(
            || {
                (
                    rand_value::<E>(),
                    rand_value::<E>(),
                    rand_value::<E>(),
                    rand_value::<E>(),
                )
            },
            |(mut a, mut b, mut c, mut d)| {
                for _ in 0..25 {
                    let t0 = a + b;
                    let t1 = b + c;
                    let t2 = c + d;
                    let t3 = d + a;

                    a = t0;
                    b = t1;
                    c = t2;
                    d = t3;
                }
                (a, b, c, d)
            },
            BatchSize::SmallInput,
        )
    });

    group.bench_function(format!("{extension}/batch/sub"), |b| {
        b.iter_batched(
            || {
                (
                    rand_value::<E>(),
                    rand_value::<E>(),
                    rand_value::<E>(),
                    rand_value::<E>(),
                )
            },
            |(mut a, mut b, mut c, mut d)| {
                for _ in 0..25 {
                    let t0 = a - b;
                    let t1 = b - c;
                    let t2 = c - d;
                    let t3 = d - a;

                    a = t0;
                    b = t1;
                    c = t2;
                    d = t3;
                }
                (a, b, c, d)
            },
            BatchSize::SmallInput,
        )
    });

    group.bench_function(format!("{extension}/batch/mul"), |b| {
        b.iter_batched(
            || {
                (
                    rand_value::<E>(),
                    rand_value::<E>(),
                    rand_value::<E>(),
                    rand_value::<E>(),
                )
            },
            |(mut a, mut b, mut c, mut d)| {
                for _ in 0..25 {
                    let t0 = a * b;
                    let t1 = b * c;
                    let t2 = c * d;
                    let t3 = d * a;

                    a = t0;
                    b = t1;
                    c = t2;
                    d = t3;
                }
                (a, b, c, d)
            },
            BatchSize::SmallInput,
        )
    });
}

// GENERIC BENCHMARK RUNNER
// ================================================================================================

fn bench_field_ops(c: &mut Criterion) {
    field_ops::<f62::BaseElement>(c, "f62");
    field_ops::<f64::BaseElement>(c, "f64");
    field_ops::<f128::BaseElement>(c, "f128");
}

// CRITERION BOILERPLATE
// ================================================================================================

criterion_group!(field_group, batch_inv, bench_field_ops);
criterion_main!(field_group);


// Copyright (c) Facebook, Inc. and its affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
use rand_utils::rand_vector;
use std::time::Duration;
use winter_math::{fft, fields::f128::BaseElement, polynom, FieldElement};

const SIZES: [usize; 3] = [262_144, 524_288, 1_048_576];

fn syn_div(c: &mut Criterion) {
    let mut group = c.benchmark_group("syn_div");
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(10));

    for &size in SIZES.iter() {
        let stride = 8;
        let mut values: Vec<BaseElement> = rand_vector(size);
        for v in values.iter_mut().skip(stride) {
            *v = BaseElement::ZERO;
        }
        let inv_twiddles = fft::get_inv_twiddles::<BaseElement>(size);
        fft::interpolate_poly(&mut values, &inv_twiddles);
        let p = values;
        let z_power = size / stride;

        group.bench_function(BenchmarkId::new("high_degree", size), |bench| {
            bench.iter_batched_ref(
                || p.clone(),
                |p| polynom::syn_div(p, z_power, BaseElement::ONE),
                BatchSize::LargeInput,
            );
        });
    }

    group.finish();
}

criterion_group!(polynom_group, syn_div);
criterion_main!(polynom_group);


// Copyright (c) Facebook, Inc. and its affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

//! This crate contains modules with mathematical operations needed in STARK proof generation
//! and verification.
//!
//! Unless otherwise stated, all operations happen in [finite fields](https://en.wikipedia.org/wiki/Finite_field).
//!
//! # Finite fields
//! [Finite field](fields) modules implements arithmetic operations in STARK-friendly finite
//! fields. The operation include:
//!
//! * Basic arithmetic operations: addition, multiplication, subtraction, division, inversion.
//! * Drawing random and pseudo-random elements from the field.
//! * Computing roots of unity of a given order.
//!
//! Currently, there are two implementations of finite fields:
//!
//! * A 128-bit field with modulus 2128 - 45 * 240 + 1. This field was not
//!   chosen with any significant thought given to performance, and the implementation of most
//!   operations is sub-optimal as well. Proofs generated in this field can support security level
//!   of ~100 bits. If higher level of security is desired, proofs must be generated in a quadratic
//!   extension of the field.
//! * A 62-bit field with modulus 262 - 111 * 239 + 1. This field supports
//!   very fast modular arithmetic including branchless multiplication and addition. To achieve
//!   adequate security (i.e. ~100 bits), proofs must be generated in a quadratic extension of this
//!   field. For higher levels of security, a cubic extension field should be used.
//! * A 64-bit field with modulus 264 - 232 + 1. This field is about 15%
//!   slower than the 62-bit field described above, but it has a number of other attractive
//!   properties. To achieve adequate security (i.e. ~100 bits), proofs must be generated in a
//!   quadratic extension of this field. For higher levels of security, a cubic extension field
//!   should be used.
//!
//! ## Extension fields
//!
//! Currently, the library provides a generic way to create quadratic and cubic extensions of
//! supported STARK fields. This can be done by implementing [ExtensibleField] trait for
//! degrees 2 and 3.
//!
//! Quadratic extension fields are defined using the following irreducible polynomials:
//! * For [f62](crate::fields::f62) field, the polynomial is x2 - x - 1.
//! * For [f64](crate::fields::f64) field, the polynomial is x2 - x + 2.
//! * For [f128](crate::fields::f128) field, the polynomial is x2 - x - 1.
//!
//! Cubic extension fields are defined using the following irreducible polynomials:
//! * For [f62](crate::fields::f62) field, the polynomial is x3 + 2x + 2.
//! * For [f64](crate::fields::f64) field, the polynomial is x3 - x - 1.
//! * For [f128](crate::fields::f128) field, cubic extensions are not supported.
//!
//! # Polynomials
//! [Polynomials](polynom) module implements basic polynomial operations such as:
//!
//! * Evaluation of a polynomial at a single or multiple point.
//! * Interpolation of a polynomial from a set of points (using
//!   [Lagrange](https://en.wikipedia.org/wiki/Lagrange_polynomial) interpolation).
//! * Addition, multiplication, subtraction, and division of polynomials.
//! * Synthetic polynomial division (using
//!   [Ruffini's](https://en.wikipedia.org/wiki/Ruffini%27s_rule) method).
//!
//! # Fast Fourier transform
//! [FFT](fft) module contains operations for computing Fast Fourier transform in a prime
//! field (also called [Number-theoretic transform](https://en.wikipedia.org/wiki/Discrete_Fourier_transform_(general)#Number-theoretic_transform)).
//! This can be used to interpolate and evaluate polynomials in *O(n log n)* time as long as
//! the domain of the polynomial is a multiplicative subgroup with size which is a power of 2.
//!
//! # Concurrent execution
//!
//! When the crate is compiled with `concurrent` feature enabled, some operations will be
//! executed in multiple threads (usually, as many threads as there are logical cores on
//! the machine). These operations are:
//!
//! * crate:
//!   - [get_power_series()]
//!   - [get_power_series_with_offset()]
//!   - [add_in_place()]
//!   - [mul_acc()]
//!   - [batch_inversion()]
//! * `fft` module:
//!   - [evaluate_poly()](fft::evaluate_poly())
//!   - [evaluate_poly_with_offset()](fft::evaluate_poly_with_offset())
//!   - [interpolate_poly()](fft::interpolate_poly())
//!   - [interpolate_poly_with_offset()][fft::interpolate_poly_with_offset()]
//!   - [get_twiddles()](fft::get_twiddles())
//!   - [get_inv_twiddles()](fft::get_twiddles())
//!
//! Number of threads can be configured via `RAYON_NUM_THREADS` environment variable

#![cfg_attr(not(feature = "std"), no_std)]

#[cfg(not(feature = "std"))]
#[macro_use]
extern crate alloc;

pub mod fft;
pub mod polynom;

mod field;
pub use field::{ExtensibleField, ExtensionOf, FieldElement, StarkField, ToElements};
pub mod fields {
    //! Finite field implementations.
    //!
    //! This module contains concrete implementations of base STARK fields as well as extensions
    //! of these field.

    pub use super::field::f128;
    pub use super::field::f62;
    pub use super::field::f64;
    pub use super::field::CubeExtension;
    pub use super::field::QuadExtension;
}

mod utils;
pub use crate::utils::{
    add_in_place, batch_inversion, get_power_series, get_power_series_with_offset, log2, mul_acc,
};