// 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::StarkDomain;
use air::{Air, AuxTraceRandElements, ConstraintDivisor};
use math::{fft, ExtensionOf, FieldElement};
use utils::collections::{BTreeMap, Vec};

// CONSTANTS
// ================================================================================================

/// Boundary polynomials with this degree or smaller will be evaluated on the fly, while for
/// larger polynomials all evaluations over the constraint evaluation domain will be pre-computed.
const SMALL_POLY_DEGREE: usize = 63;

// BOUNDARY CONSTRAINTS
// ================================================================================================

/// Contains all boundary constraints defined for an instance of a computation. This includes
/// constraints against the main segment of the execution trace as well as constraints against
/// auxiliary trace segments (if any).
///
/// We transform the constraints defined in the [air] crate into specialized constraints here
/// to make evaluation of these constraints more efficient in the prover context.
pub struct BoundaryConstraints<E: FieldElement>(Vec<BoundaryConstraintGroup<E>>);

impl<E: FieldElement> BoundaryConstraints<E> {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new instance of [BoundaryConstraints] constructed from the constraints defined
    /// by an instance of AIR for a specific computation.
    pub fn new<A: Air<BaseField = E::BaseField>>(
        air: &A,
        aux_rand_elements: &AuxTraceRandElements<E>,
        composition_coefficients: &[E],
    ) -> Self {
        // get constraints from the AIR instance
        let source = air.get_boundary_constraints(aux_rand_elements, composition_coefficients);

        // initialize a map of twiddles here so that we can keep track of already computed
        // twiddles; this helps us avoid building twiddles over and over again for constraints
        // defined over the same domain. twiddles are relevant only for large polynomial
        // constraints.
        let mut twiddle_map = BTreeMap::new();

        // transform constraints against the main segment of the execution trace into specialized
        // constraints
        let mut result = source
            .main_constraints()
            .iter()
            .map(|group| {
                BoundaryConstraintGroup::from_main_constraints(group, air, &mut twiddle_map)
            })
            .collect::<Vec<BoundaryConstraintGroup<E>>>();

        // transform constraints against auxiliary trace segments (if any) into specialized
        // constraints. this also checks if a group with the same divisor has already been
        // transformed (when processing constraints against the main trace above), and if so,
        // appends constraints to that group rather than creating a new group. this ensures
        // that we always end up with a single constraint group for the same divisor.
        for group in source.aux_constraints() {
            match result.iter_mut().find(|g| &g.divisor == group.divisor()) {
                Some(x) => x.add_aux_constraints(group, air, &mut twiddle_map),
                None => {
                    let group =
                        BoundaryConstraintGroup::from_aux_constraints(group, air, &mut twiddle_map);
                    result.push(group);
                }
            };
        }

        Self(result)
    }

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

    /// Returns a vector of all boundary constraint divisors.
    pub fn get_divisors(&self) -> Vec<ConstraintDivisor<E::BaseField>> {
        self.0.iter().map(|g| g.divisor.clone()).collect()
    }

    // EVALUATORS
    // --------------------------------------------------------------------------------------------

    /// Evaluates boundary constraints against the main segment of an execution trace at the
    /// specified step of constraint evaluation domain.
    pub fn evaluate_main(
        &self,
        main_state: &[E::BaseField],
        domain: &StarkDomain<E::BaseField>,
        step: usize,
        result: &mut [E],
    ) {
        let x = domain.get_ce_x_at(step);
        for (group, result) in self.0.iter().zip(result.iter_mut()) {
            *result = group.evaluate_main(main_state, step, x);
        }
    }

    /// Evaluates boundary constraints against all segments of an execution trace at the
    /// specified step of constraint evaluation domain.
    pub fn evaluate_all(
        &self,
        main_state: &[E::BaseField],
        aux_state: &[E],
        domain: &StarkDomain<E::BaseField>,
        step: usize,
        result: &mut [E],
    ) {
        let x = domain.get_ce_x_at(step);
        for (group, result) in self.0.iter().zip(result.iter_mut()) {
            *result = group.evaluate_all(main_state, aux_state, step, x);
        }
    }
}

// BOUNDARY CONSTRAINT GROUP
// ================================================================================================

/// Contains constraints all having the same divisor. The constraints are separated into single
/// value constraints, small polynomial constraints, and large polynomial constraints.
///
/// The constraints are also separated into constraints against the main segment of the execution
/// and the constraints against auxiliary segments of the execution trace (if any).
pub struct BoundaryConstraintGroup<E: FieldElement> {
    divisor: ConstraintDivisor<E::BaseField>,
    // main trace constraints
    main_single_value: Vec<SingleValueConstraint<E::BaseField, E>>,
    main_small_poly: Vec<SmallPolyConstraint<E::BaseField, E>>,
    main_large_poly: Vec<LargePolyConstraint<E::BaseField, E>>,
    // auxiliary trace constraints
    aux_single_value: Vec<SingleValueConstraint<E, E>>,
    aux_small_poly: Vec<SmallPolyConstraint<E, E>>,
    aux_large_poly: Vec<LargePolyConstraint<E, E>>,
}

impl<E: FieldElement> BoundaryConstraintGroup<E> {
    // CONSTRUCTORS
    // --------------------------------------------------------------------------------------------

    /// Returns an empty [BoundaryConstraintGroup] instantiated with the specified divisor and
    /// degree adjustment factor.
    fn new(divisor: ConstraintDivisor<E::BaseField>) -> Self {
        Self {
            divisor,
            main_single_value: Vec::new(),
            main_small_poly: Vec::new(),
            main_large_poly: Vec::new(),
            aux_single_value: Vec::new(),
            aux_small_poly: Vec::new(),
            aux_large_poly: Vec::new(),
        }
    }

    /// Returns a [BoundaryConstraintGroup] created from the specified group of constraints against
    /// the main segment of an execution trace. Constraints against auxiliary trace segment in this
    /// group will be empty.
    ///
    /// Twiddles and [Air] instance are passed in for evaluating large polynomial constraints
    /// (if any).
    pub fn from_main_constraints<A: Air<BaseField = E::BaseField>>(
        source: &air::BoundaryConstraintGroup<E::BaseField, E>,
        air: &A,
        twiddle_map: &mut BTreeMap<usize, Vec<E::BaseField>>,
    ) -> Self {
        let mut result = Self::new(source.divisor().clone());

        for constraint in source.constraints() {
            if constraint.poly().len() == 1 {
                let constraint = SingleValueConstraint::new(constraint);
                result.main_single_value.push(constraint);
            } else if constraint.poly().len() < SMALL_POLY_DEGREE {
                let constraint = SmallPolyConstraint::new(constraint);
                result.main_small_poly.push(constraint);
            } else {
                let constraint = LargePolyConstraint::new(constraint, air, twiddle_map);
                result.main_large_poly.push(constraint);
            }
        }

        result
    }

    /// Returns a [BoundaryConstraintGroup] created from the specified group of constraints against
    /// auxiliary segments of an execution trace. Constraints against the main trace segment in this
    /// group will be empty.
    ///
    /// Twiddles and [Air] instance are passed in for evaluating large polynomial constraints
    /// (if any).
    pub fn from_aux_constraints<A: Air<BaseField = E::BaseField>>(
        group: &air::BoundaryConstraintGroup<E, E>,
        air: &A,
        twiddle_map: &mut BTreeMap<usize, Vec<E::BaseField>>,
    ) -> Self {
        let mut result = Self::new(group.divisor().clone());
        result.add_aux_constraints(group, air, twiddle_map);
        result
    }

    // STATE MUTATORS
    // --------------------------------------------------------------------------------------------

    /// Adds the provided constraints against auxiliary segments of an execution trace to this
    /// group.
    ///
    /// Twiddles and [Air] instance are passed in for evaluating large polynomial constraints
    /// (if any).
    ///
    /// # Panics
    /// Panics if the divisor of the provided constraints doesn't match the divisor of this group.
    pub fn add_aux_constraints<A: Air<BaseField = E::BaseField>>(
        &mut self,
        group: &air::BoundaryConstraintGroup<E, E>,
        air: &A,
        twiddle_map: &mut BTreeMap<usize, Vec<E::BaseField>>,
    ) {
        assert_eq!(
            group.divisor(),
            &self.divisor,
            "inconsistent constraint divisor"
        );

        for constraint in group.constraints() {
            if constraint.poly().len() == 1 {
                let constraint = SingleValueConstraint::new(constraint);
                self.aux_single_value.push(constraint);
            } else if constraint.poly().len() < SMALL_POLY_DEGREE {
                let constraint = SmallPolyConstraint::new(constraint);
                self.aux_small_poly.push(constraint);
            } else {
                let constraint = LargePolyConstraint::new(constraint, air, twiddle_map);
                self.aux_large_poly.push(constraint);
            }
        }
    }

    // EVALUATORS
    // --------------------------------------------------------------------------------------------

    /// Evaluates the constraints against the main segment of the execution trace contained in
    /// this group at the specified step of the trace.
    pub fn evaluate_main(&self, state: &[E::BaseField], ce_step: usize, x: E::BaseField) -> E {
        let mut result = E::ZERO;

        // evaluate all single-value constraints
        for constraint in self.main_single_value.iter() {
            result += constraint.evaluate(state);
        }

        // evaluate all small polynomial constraints
        for constraint in self.main_small_poly.iter() {
            result += constraint.evaluate(state, x);
        }

        // evaluate all large polynomial constraints
        for constraint in self.main_large_poly.iter() {
            result += constraint.evaluate(state, ce_step);
        }

        result
    }

    /// Evaluates all constraints contained in this group at the specified step of the
    /// execution trace.
    pub fn evaluate_all(
        &self,
        main_state: &[E::BaseField],
        aux_state: &[E],
        ce_step: usize,
        x: E::BaseField,
    ) -> E {
        let mut result = self.evaluate_main(main_state, ce_step, x);

        // evaluate all single-value constraints
        for constraint in self.aux_single_value.iter() {
            result += constraint.evaluate(aux_state);
        }

        // evaluate all small polynomial constraints
        for constraint in self.aux_small_poly.iter() {
            result += constraint.evaluate(aux_state, x);
        }

        // evaluate all large polynomial constraints
        for constraint in self.aux_large_poly.iter() {
            result += constraint.evaluate(aux_state, ce_step);
        }

        result
    }
}

// CONSTRAINT SPECIALIZATIONS
// ================================================================================================

/// A constraint where the numerator can be represented by p(x) - v, where v is the asserted value,
/// and p(x) is the trace polynomial for the column against which the constraint is applied.
struct SingleValueConstraint<F, E>
where
    F: FieldElement,
    E: FieldElement<BaseField = F::BaseField> + ExtensionOf<F>,
{
    column: usize,
    value: F,
    coefficients: E,
}

impl<F, E> SingleValueConstraint<F, E>
where
    F: FieldElement,
    E: FieldElement<BaseField = F::BaseField> + ExtensionOf<F>,
{
    /// Returns an new instance of [SingleValueConstraint] created from the specified source
    /// boundary constraint.
    pub fn new(source: &air::BoundaryConstraint<F, E>) -> Self {
        debug_assert!(source.poly().len() == 1, "not a single constraint");
        Self {
            column: source.column(),
            value: source.poly()[0],
            coefficients: *source.cc(),
        }
    }

    /// Evaluates this constraint over the specified state and returns the result.
    ///
    /// This also multiplies by the composition coefficient.
    pub fn evaluate(&self, state: &[F]) -> E {
        let evaluation = state[self.column] - self.value;
        self.coefficients.mul_base(evaluation)
    }
}

/// A constraint where the numerator can be represented by p(x) - c(x), where b(x) is the
/// polynomial describing a set of asserted values. This specialization is useful when the
/// degree of b(x) is relatively small, and thus, is cheap to evaluate on the fly.
///
/// TODO: investigate whether we get any significant improvement vs. [LargePolyConstraint], and if
/// so, what is the appropriate value for SMALL_POLY_DEGREE.
struct SmallPolyConstraint<F, E>
where
    F: FieldElement,
    E: FieldElement<BaseField = F::BaseField> + ExtensionOf<F>,
{
    column: usize,
    poly: Vec<F>,
    x_offset: F::BaseField,
    coefficients: E,
}

impl<F, E> SmallPolyConstraint<F, E>
where
    F: FieldElement,
    E: FieldElement<BaseField = F::BaseField> + ExtensionOf<F>,
{
    /// Returns an new instance of [SmallPolyConstraint] created from the specified source
    /// boundary constraint.
    pub fn new(source: &air::BoundaryConstraint<F, E>) -> Self {
        debug_assert!(
            source.poly().len() > 1 && source.poly().len() < SMALL_POLY_DEGREE,
            "not a small poly constraint"
        );
        Self {
            column: source.column(),
            poly: source.poly().to_vec(),
            x_offset: source.poly_offset().1,
            coefficients: *source.cc(),
        }
    }

    /// Evaluates this constraint over the specified state and returns the result.
    ///
    /// We need the point in the domain ('x') corresponding to the passed-in state because to
    /// evaluate the constraint we need to evaluate its value polynomial at `x`.
    pub fn evaluate(&self, state: &[F], x: F::BaseField) -> E {
        let x = x * self.x_offset;
        // evaluate constraint polynomial as x * offset using Horner evaluation
        let assertion_value = self
            .poly
            .iter()
            .rev()
            .fold(F::ZERO, |acc, &coeff| acc.mul_base(x) + coeff);
        // evaluate the constraint
        let evaluation = state[self.column] - assertion_value;
        self.coefficients.mul_base(evaluation)
    }
}

/// A constraint where the numerator can be represented by p(x) - b(x), where b(x) is a large
/// polynomial. In such cases, we pre-compute evaluations of b(x) by evaluating it over the
/// entire constraint evaluation domain (using FFT).
struct LargePolyConstraint<F, E>
where
    F: FieldElement,
    E: FieldElement<BaseField = F::BaseField> + ExtensionOf<F>,
{
    column: usize,
    values: Vec<F>,
    step_offset: usize,
    coefficients: E,
}

impl<F, E> LargePolyConstraint<F, E>
where
    F: FieldElement,
    E: FieldElement<BaseField = F::BaseField> + ExtensionOf<F>,
{
    /// Returns a new instance of [LargePolyConstraint] created from the specified source
    /// boundary constraint.
    pub fn new<A: Air<BaseField = F::BaseField>>(
        source: &air::BoundaryConstraint<F, E>,
        air: &A,
        twiddle_map: &mut BTreeMap<usize, Vec<F::BaseField>>,
    ) -> Self {
        debug_assert!(
            source.poly().len() >= SMALL_POLY_DEGREE,
            "not a large poly constraint"
        );
        // evaluate the polynomial over the entire constraint evaluation domain; first
        // get twiddles for the evaluation; if twiddles haven't been built yet, build them
        let poly_length = source.poly().len();
        let twiddles = twiddle_map
            .entry(poly_length)
            .or_insert_with(|| fft::get_twiddles(poly_length));

        let values = fft::evaluate_poly_with_offset(
            source.poly(),
            twiddles,
            air.domain_offset(),
            air.ce_domain_size() / poly_length,
        );

        LargePolyConstraint {
            column: source.column(),
            values,
            step_offset: source.poly_offset().0 * air.ce_blowup_factor(),
            coefficients: *source.cc(),
        }
    }

    /// Evaluates this constraint at the specified step of the constraint evaluation domain.
    ///
    /// This also applies composition coefficients as well as the degree adjustment factor
    /// (defined by `xp` parameter) to the evaluation before it is returned.
    pub fn evaluate(&self, state: &[F], ce_step: usize) -> E {
        let value_index = if self.step_offset > 0 {
            // if the assertion happens on steps which are not a power of 2, we need to offset the
            // evaluation; the below basically computes (ce_step - step_offset) % values.len();
            // this is equivalent to evaluating the polynomial at x * x_offset coordinate.
            if self.step_offset > ce_step {
                self.values.len() + ce_step - self.step_offset
            } else {
                ce_step - self.step_offset
            }
        } else {
            ce_step
        };
        let evaluation = state[self.column] - self.values[value_index];
        (self.coefficients).mul_base(evaluation)
    }
}


// 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::RowMatrix;
use air::proof::Queries;
use crypto::{ElementHasher, MerkleTree};
use math::FieldElement;
use utils::collections::Vec;

// CONSTRAINT COMMITMENT
// ================================================================================================

/// Constraint evaluation commitment.
///
/// The commitment consists of two components:
/// * Evaluations of composition polynomial columns over the LDE domain.
/// * Merkle tree where each leaf in the tree corresponds to a row in the composition polynomial
///   evaluation matrix.
pub struct ConstraintCommitment<E: FieldElement, H: ElementHasher<BaseField = E::BaseField>> {
    evaluations: RowMatrix<E>,
    commitment: MerkleTree<H>,
}

impl<E: FieldElement, H: ElementHasher<BaseField = E::BaseField>> ConstraintCommitment<E, H> {
    /// Creates a new constraint evaluation commitment from the provided composition polynomial
    /// evaluations and the corresponding Merkle tree commitment.
    pub fn new(evaluations: RowMatrix<E>, commitment: MerkleTree<H>) -> ConstraintCommitment<E, H> {
        assert_eq!(
            evaluations.num_rows(),
            commitment.leaves().len(),
            "number of rows in constraint evaluation matrix must be the same as number of leaves in constraint commitment"
        );
        ConstraintCommitment {
            evaluations,
            commitment,
        }
    }

    /// Returns the root of the commitment Merkle tree.
    pub fn root(&self) -> H::Digest {
        *self.commitment.root()
    }

    /// Returns the depth of the commitment Merkle tree.
    #[allow(unused)]
    pub fn tree_depth(&self) -> usize {
        self.commitment.depth()
    }

    /// Returns constraint evaluations at the specified positions along with Merkle authentication
    /// paths from the root of the commitment to these evaluations.
    pub fn query(self, positions: &[usize]) -> Queries {
        // build Merkle authentication paths to the leaves specified by positions
        let merkle_proof = self
            .commitment
            .prove_batch(positions)
            .expect("failed to generate a Merkle proof for constraint queries");

        // determine a set of evaluations corresponding to each position
        let mut evaluations = Vec::new();
        for &position in positions {
            let row = self.evaluations.row(position).to_vec();
            evaluations.push(row);
        }

        Queries::new(merkle_proof, evaluations)
    }
}


// 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::ColMatrix;
use math::{polynom::degree_of, FieldElement};
use utils::collections::Vec;

// COMPOSITION POLYNOMIAL
// ================================================================================================
/// Represents a composition polynomial split into columns with each column being of length equal
/// to trace_length. Thus, for example, if the composition polynomial has degree 2N - 1, where N
/// is the trace length, it will be stored as two columns of size N (each of degree N - 1).
pub struct CompositionPoly<E: FieldElement> {
    data: ColMatrix<E>,
}

impl<E: FieldElement> CompositionPoly<E> {
    /// Returns a new composition polynomial.
    pub fn new(coefficients: Vec<E>, trace_length: usize, num_cols: usize) -> Self {
        assert!(
            coefficients.len().is_power_of_two(),
            "size of composition polynomial must be a power of 2, but was {}",
            coefficients.len(),
        );
        assert!(
            trace_length.is_power_of_two(),
            "trace length must be a power of 2, but was {trace_length}"
        );
        assert!(
            trace_length < coefficients.len(),
            "trace length must be smaller than size of composition polynomial"
        );

        let polys = segment(coefficients, trace_length, num_cols);

        CompositionPoly {
            data: ColMatrix::new(polys),
        }
    }

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

    /// Returns the number of individual column polynomials used to describe this composition
    /// polynomial.
    pub fn num_columns(&self) -> usize {
        self.data.num_cols()
    }

    /// Returns the length of individual column polynomials; this is guaranteed to be a power of 2.
    pub fn column_len(&self) -> usize {
        self.data.num_rows()
    }

    /// Returns the degree of individual column polynomial.
    #[allow(unused)]
    pub fn column_degree(&self) -> usize {
        self.column_len() - 1
    }

    /// Returns evaluations of all composition polynomial columns at point z.
    pub fn evaluate_at(&self, z: E) -> Vec<E> {
        self.data.evaluate_columns_at(z)
    }

    /// Returns a reference to the matrix of individual column polynomials.
    pub fn data(&self) -> &ColMatrix<E> {
        &self.data
    }

    /// Transforms this composition polynomial into a vector of individual column polynomials.
    pub fn into_columns(self) -> Vec<Vec<E>> {
        self.data.into_columns()
    }
}

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

/// Splits polynomial coefficients into the specified number of columns. The coefficients are split
/// in such a way that each resulting column has the same degree. For example, a polynomial
/// a * x^3 + b * x^2 + c * x + d, can be rewritten as: (c * x + d) + x^2 * (a * x + b), and then
/// the two columns will be: (c * x + d) and (a * x + b).
fn segment<E: FieldElement>(
    coefficients: Vec<E>,
    trace_len: usize,
    num_cols: usize,
) -> Vec<Vec<E>> {
    debug_assert!(degree_of(&coefficients) < trace_len * num_cols);

    coefficients
        .chunks(trace_len)
        .take(num_cols)
        .map(|slice| slice.to_vec())
        .collect()
}

// TESTS
// ================================================================================================

#[cfg(test)]
mod tests {

    use math::fields::f128::BaseElement;
    use utils::collections::Vec;

    #[test]
    fn segment() {
        let values = (0u128..16).map(BaseElement::new).collect::<Vec<_>>();
        let actual = super::segment(values, 4, 4);

        #[rustfmt::skip]
        let expected = vec![
            vec![BaseElement::new(0), BaseElement::new(1), BaseElement::new(2), BaseElement::new(3)],
            vec![BaseElement::new(4), BaseElement::new(5), BaseElement::new(6), BaseElement::new(7)],
            vec![BaseElement::new(8), BaseElement::new(9), BaseElement::new(10), BaseElement::new(11)],
            vec![BaseElement::new(12), BaseElement::new(13), BaseElement::new(14), BaseElement::new(15)],
        ];

        assert_eq!(expected, actual)
    }
}


// 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::{CompositionPoly, ConstraintDivisor, ProverError, StarkDomain};
use math::{batch_inversion, fft, FieldElement, StarkField};
use utils::{batch_iter_mut, collections::Vec, iter_mut, uninit_vector};

#[cfg(debug_assertions)]
use air::TransitionConstraints;

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

// CONSTANTS
// ================================================================================================

const MIN_FRAGMENT_SIZE: usize = 16;

// CONSTRAINT EVALUATION TABLE
// ================================================================================================

pub struct ConstraintEvaluationTable<'a, E: FieldElement> {
    evaluations: Vec<Vec<E>>,
    divisors: Vec<ConstraintDivisor<E::BaseField>>,
    domain: &'a StarkDomain<E::BaseField>,

    #[cfg(debug_assertions)]
    main_transition_evaluations: Vec<Vec<E::BaseField>>,
    #[cfg(debug_assertions)]
    aux_transition_evaluations: Vec<Vec<E>>,
    #[cfg(debug_assertions)]
    expected_transition_degrees: Vec<usize>,
}

impl<'a, E: FieldElement> ConstraintEvaluationTable<'a, E> {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new constraint evaluation table with number of columns equal to the number of
    /// specified divisors, and number of rows equal to the size of constraint evaluation domain.
    #[cfg(not(debug_assertions))]
    pub fn new(
        domain: &'a StarkDomain<E::BaseField>,
        divisors: Vec<ConstraintDivisor<E::BaseField>>,
    ) -> Self {
        let num_columns = divisors.len();
        let num_rows = domain.ce_domain_size();
        ConstraintEvaluationTable {
            evaluations: uninit_matrix(num_columns, num_rows),
            divisors,
            domain,
        }
    }

    /// Similar to the as above constructor but used in debug mode. In debug mode we also want
    /// to keep track of all evaluated transition constraints so that we can verify that their
    /// expected degrees match their actual degrees.
    #[cfg(debug_assertions)]
    pub fn new(
        domain: &'a StarkDomain<E::BaseField>,
        divisors: Vec<ConstraintDivisor<E::BaseField>>,
        transition_constraints: &TransitionConstraints<E>,
    ) -> Self {
        let num_columns = divisors.len();
        let num_rows = domain.ce_domain_size();
        let num_tm_columns = transition_constraints.num_main_constraints();
        let num_ta_columns = transition_constraints.num_aux_constraints();

        // collect expected degrees for all transition constraints to compare them against actual
        // degrees; we do this in debug mode only because this comparison is expensive
        let expected_transition_degrees =
            build_transition_constraint_degrees(transition_constraints, domain.trace_length());

        ConstraintEvaluationTable {
            evaluations: uninit_matrix(num_columns, num_rows),
            divisors,
            domain,
            main_transition_evaluations: uninit_matrix(num_tm_columns, num_rows),
            aux_transition_evaluations: uninit_matrix(num_ta_columns, num_rows),
            expected_transition_degrees,
        }
    }

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

    /// Returns the number of rows in this table. This is the same as the size of the constraint
    /// evaluation domain.
    pub fn num_rows(&self) -> usize {
        self.evaluations[0].len()
    }

    /// Returns number of columns in this table. The first column always contains the value of
    /// combined transition constraint evaluations; the remaining columns contain values of
    /// assertion constraint evaluations combined based on common divisors.
    #[allow(dead_code)]
    pub fn num_columns(&self) -> usize {
        self.evaluations.len()
    }

    // TABLE FRAGMENTS
    // --------------------------------------------------------------------------------------------

    /// Break the table into the number of specified fragments. All fragments can be updated
    /// independently - e.g. in different threads.
    pub fn fragments(&mut self, num_fragments: usize) -> Vec<EvaluationTableFragment<E>> {
        let fragment_size = self.num_rows() / num_fragments;
        assert!(
            fragment_size >= MIN_FRAGMENT_SIZE,
            "fragment size must be at least {MIN_FRAGMENT_SIZE}, but was {fragment_size}"
        );

        // break evaluations into fragments
        let evaluation_data = make_fragments(&mut self.evaluations, num_fragments);

        #[cfg(debug_assertions)]
        let result = {
            // in debug mode, also break individual transition evaluations into fragments
            let tm_evaluation_data =
                make_fragments(&mut self.main_transition_evaluations, num_fragments);
            let ta_evaluation_data =
                make_fragments(&mut self.aux_transition_evaluations, num_fragments);

            evaluation_data
                .into_iter()
                .zip(tm_evaluation_data)
                .zip(ta_evaluation_data)
                .enumerate()
                .map(|(i, ((evaluations, tm_evaluations), ta_evaluations))| {
                    EvaluationTableFragment {
                        offset: i * fragment_size,
                        evaluations,
                        tm_evaluations,
                        ta_evaluations,
                    }
                })
                .collect()
        };

        #[cfg(not(debug_assertions))]
        let result = evaluation_data
            .into_iter()
            .enumerate()
            .map(|(i, evaluations)| EvaluationTableFragment {
                offset: i * fragment_size,
                evaluations,
            })
            .collect();

        result
    }

    // CONSTRAINT COMPOSITION
    // --------------------------------------------------------------------------------------------
    /// Divides constraint evaluation columns by their respective divisor (in evaluation form),
    /// combines the results into a single column, and interpolates this column into a composition
    /// polynomial in coefficient form.
    /// `num_cols` is the number of necessary columns (of length `trace_length`) needed to store
    /// the coefficients of the constraint composition polynomial and is needed by
    /// `CompositionPoly::new`.
    pub fn into_poly(self, num_cols: usize) -> Result<CompositionPoly<E>, ProverError> {
        // allocate memory for the combined polynomial
        let mut combined_poly = E::zeroed_vector(self.num_rows());

        // iterate over all columns of the constraint evaluation table, divide each column
        // by the evaluations of its corresponding divisor, and add all resulting evaluations
        // together into a single vector
        for (column, divisor) in self.evaluations.into_iter().zip(self.divisors.iter()) {
            // divide the column by the divisor and accumulate the result into combined_poly
            acc_column(column, divisor, self.domain, &mut combined_poly);
        }

        // at this point, combined_poly contains evaluations of the combined constraint polynomial;
        // we interpolate this polynomial to transform it into coefficient form.
        let inv_twiddles = fft::get_inv_twiddles::<E::BaseField>(combined_poly.len());
        fft::interpolate_poly_with_offset(&mut combined_poly, &inv_twiddles, self.domain.offset());

        let trace_length = self.domain.trace_length();
        Ok(CompositionPoly::new(combined_poly, trace_length, num_cols))
    }

    // DEBUG HELPERS
    // --------------------------------------------------------------------------------------------

    #[cfg(debug_assertions)]
    pub fn validate_transition_degrees(&mut self) {
        // evaluate transition constraint divisor (which is assumed to be the first one in the
        // divisor list) over the constraint evaluation domain. this is used later to compute
        // actual degrees of transition constraint evaluations.
        let div_values = evaluate_divisor::<E::BaseField>(
            &self.divisors[0],
            self.num_rows(),
            self.domain.offset(),
        );

        // collect actual degrees for all transition constraints by interpolating saved
        // constraint evaluations into polynomials and checking their degree; also
        // determine max transition constraint degree
        let mut actual_degrees = Vec::with_capacity(self.expected_transition_degrees.len());
        let mut max_degree = 0;
        let inv_twiddles = fft::get_inv_twiddles::<E::BaseField>(self.num_rows());

        // first process transition constraint evaluations for the main trace segment
        for evaluations in self.main_transition_evaluations.iter() {
            let degree = get_transition_poly_degree(evaluations, &inv_twiddles, &div_values);
            actual_degrees.push(degree);
            max_degree = core::cmp::max(max_degree, degree);
        }

        // then process transition constraint evaluations for auxiliary trace segments
        for evaluations in self.aux_transition_evaluations.iter() {
            let degree = get_transition_poly_degree(evaluations, &inv_twiddles, &div_values);
            actual_degrees.push(degree);
            max_degree = core::cmp::max(max_degree, degree);
        }

        // make sure expected and actual degrees are equal
        assert_eq!(
            self.expected_transition_degrees, actual_degrees,
            "transition constraint degrees didn't match\nexpected: {:>3?}\nactual:   {:>3?}",
            self.expected_transition_degrees, actual_degrees
        );

        // make sure evaluation domain size does not exceed the size required by max degree
        let expected_domain_size =
            core::cmp::max(max_degree, self.domain.trace_length() + 1).next_power_of_two();
        assert_eq!(
            expected_domain_size,
            self.num_rows(),
            "incorrect constraint evaluation domain size; expected {}, but was {}",
            expected_domain_size,
            self.num_rows()
        );
    }
}

// TABLE FRAGMENTS
// ================================================================================================

pub struct EvaluationTableFragment<'a, E: FieldElement> {
    offset: usize,
    evaluations: Vec<&'a mut [E]>,

    #[cfg(debug_assertions)]
    tm_evaluations: Vec<&'a mut [E::BaseField]>,
    #[cfg(debug_assertions)]
    ta_evaluations: Vec<&'a mut [E]>,
}

impl<'a, E: FieldElement> EvaluationTableFragment<'a, E> {
    /// Returns the row at which the fragment starts.
    pub fn offset(&self) -> usize {
        self.offset
    }

    /// Returns the number of evaluation rows in the fragment.
    pub fn num_rows(&self) -> usize {
        self.evaluations[0].len()
    }

    /// Returns the number of columns in every evaluation row.
    pub fn num_columns(&self) -> usize {
        self.evaluations.len()
    }

    /// Updates a single row in the fragment with provided data.
    pub fn update_row(&mut self, row_idx: usize, row_data: &[E]) {
        for (column, &value) in self.evaluations.iter_mut().zip(row_data) {
            column[row_idx] = value;
        }
    }

    /// Updates transition evaluations row with the provided data; available only in debug mode.
    #[cfg(debug_assertions)]
    pub fn update_transition_evaluations(
        &mut self,
        row_idx: usize,
        main_evaluations: &[E::BaseField],
        aux_evaluations: &[E],
    ) {
        for (column, &value) in self.tm_evaluations.iter_mut().zip(main_evaluations) {
            column[row_idx] = value;
        }
        for (column, &value) in self.ta_evaluations.iter_mut().zip(aux_evaluations) {
            column[row_idx] = value;
        }
    }
}

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

/// Allocates memory for a two-dimensional data structure without initializing it.
fn uninit_matrix<E: FieldElement>(num_cols: usize, num_rows: usize) -> Vec<Vec<E>> {
    unsafe { (0..num_cols).map(|_| uninit_vector(num_rows)).collect() }
}

/// Breaks the source data into a mutable set of fragments such that each fragment has the same
/// number of columns as the source data, and the number of rows equal to `num_fragments`
/// parameter.
///
/// If the source data is empty, the returned vector will contain number of empty vectors equal
/// to `num_fragments` parameter.
fn make_fragments<E: FieldElement>(
    source: &mut [Vec<E>],
    num_fragments: usize,
) -> Vec<Vec<&mut [E]>> {
    let mut result = (0..num_fragments).map(|_| Vec::new()).collect::<Vec<_>>();
    if !source.is_empty() {
        let fragment_size = source[0].len() / num_fragments;
        source.iter_mut().for_each(|column| {
            for (i, fragment) in column.chunks_mut(fragment_size).enumerate() {
                result[i].push(fragment);
            }
        });
    }
    result
}

#[allow(clippy::many_single_char_names)]
fn acc_column<E: FieldElement>(
    column: Vec<E>,
    divisor: &ConstraintDivisor<E::BaseField>,
    domain: &StarkDomain<E::BaseField>,
    result: &mut [E],
) {
    let numerator = divisor.numerator();
    assert_eq!(numerator.len(), 1, "complex divisors are not yet supported");

    // compute inverse evaluations of the divisor's numerator, which has the form (x^a - b)
    let z = get_inv_evaluation(divisor, domain);

    // divide column values by the divisor; for boundary constraints this computed simply as
    // multiplication of column value by the inverse of divisor numerator; for transition
    // constraints, it is computed similarly, but the result is also multiplied by the divisor's
    // denominator (exclusion point).
    if divisor.exemptions().is_empty() {
        // the column represents merged evaluations of boundary constraints, and divisor has the
        // form of (x^a - b); thus to divide the column by the divisor, we compute: value * z,
        // where z = 1 / (x^a - 1) and has already been computed above.
        iter_mut!(result, 1024)
            .zip(column)
            .enumerate()
            .for_each(|(i, (acc_value, value))| {
                // determine which value of z corresponds to the current domain point
                let z = z[i % z.len()];
                // compute value * z and add it to the result
                *acc_value += value.mul_base(z);
            });
    } else {
        // the column represents merged evaluations of transition constraints, and divisor has the
        // form of (x^a - 1) / e(x), where e(x) describes the exemption points; thus, to divide
        // the column by the divisor, we compute: value * e(x) * z, where z = 1 / (x^a - 1) and has
        // already been computed above.
        batch_iter_mut!(
            result,
            128, // min batch size
            |batch: &mut [E], batch_offset: usize| {
                for (i, acc_value) in batch.iter_mut().enumerate() {
                    // compute value of e(x) and compute next value of x
                    let x = domain.get_ce_x_at(batch_offset + i);
                    let e = divisor.evaluate_exemptions_at(x);
                    // determine which value of z corresponds to the current domain point
                    let z = z[i % z.len()];
                    // compute value * e(x) * z and add it to the result
                    *acc_value += column[batch_offset + i].mul_base(z * e);
                }
            }
        );
    }
}

/// Computes evaluations of the divisor's numerator over the domain of the specified size and offset.
fn get_inv_evaluation<B: StarkField>(
    divisor: &ConstraintDivisor<B>,
    domain: &StarkDomain<B>,
) -> Vec<B> {
    let numerator = divisor.numerator();
    let a = numerator[0].0 as u64; // numerator degree
    let b = numerator[0].1;

    // this guarantees that we can use get_ce_x_power_at() below but limits execution trace length
    // to be at most 2^32. in the future, we should revisit this to allow execution traces of
    // greater length.
    assert!(
        a <= u32::MAX as u64,
        "constraint divisor numerator degree cannot exceed {}, but was {}",
        u32::MAX,
        a
    );

    let n = domain.ce_domain_size() / a as usize;
    let domain_offset_exp = domain.offset().exp(a.into());

    // compute x^a - b for all x
    let mut evaluations = unsafe { uninit_vector(n) };
    batch_iter_mut!(
        &mut evaluations,
        128, // min batch size
        |batch: &mut [B], batch_offset: usize| {
            for (i, evaluation) in batch.iter_mut().enumerate() {
                let x = domain.get_ce_x_power_at(batch_offset + i, a, domain_offset_exp);
                *evaluation = x - b;
            }
        }
    );

    // compute 1 / (x^a - b)
    batch_inversion(&evaluations)
}

// DEBUG HELPERS
// ================================================================================================

/// Returns evaluation degrees of all transition constraints.
///
/// An evaluation degree is defined as degree of transition constraints in the context of a given
/// execution trace accounting for constraint divisor degree. For most constraints, this degree is
/// computed as `([trace_length - 1] * [constraint degree]) - [divisor degree]`. However, for
/// constraints which rely on periodic columns this computation is slightly more complex.
///
/// The general idea is that evaluation degree is the degree of rational function `C(x) / z(x)`,
/// where `C(x)` is the constraint polynomial and `z(x)` is the divisor polynomial.
#[cfg(debug_assertions)]
fn build_transition_constraint_degrees<E: FieldElement>(
    constraints: &TransitionConstraints<E>,
    trace_length: usize,
) -> Vec<usize> {
    let mut result = Vec::new();

    for degree in constraints.main_constraint_degrees() {
        result.push(degree.get_evaluation_degree(trace_length) - constraints.divisor().degree())
    }

    for degree in constraints.aux_constraint_degrees() {
        result.push(degree.get_evaluation_degree(trace_length) - constraints.divisor().degree())
    }

    result
}

/// Computes the actual degree of a transition polynomial described by the provided evaluations.
///
/// The degree is computed as follows:
/// - First, we divide the polynomial evaluations by the evaluations of transition constraint
///   divisor (`div_values`). This is needed because it is possible for the numerator portions of
///   transition constraints to have a degree which is larger than the size of the evaluation
///   domain (and thus, interpolating the numerator would yield an incorrect result). However,
///   once the divisor values are divided out, the degree of the resulting polynomial should be
///   smaller than the size of the evaluation domain, and thus, we can interpolate safely.
/// - Then, we interpolate the polynomial over the domain specified by `inv_twiddles`.
/// - And finally, we get the degree from the interpolated polynomial.
#[cfg(debug_assertions)]
fn get_transition_poly_degree<E: FieldElement>(
    evaluations: &[E],
    inv_twiddles: &[E::BaseField],
    div_values: &[E::BaseField],
) -> usize {
    let mut evaluations = evaluations
        .iter()
        .zip(div_values)
        .map(|(&c, &d)| c / E::from(d))
        .collect::<Vec<_>>();
    fft::interpolate_poly(&mut evaluations, inv_twiddles);
    math::polynom::degree_of(&evaluations)
}

/// Evaluates constraint divisor over the specified domain. This is similar to [get_inv_evaluation]
/// function above but uses a more straight-forward but less efficient evaluation methodology and
/// also does not invert the results.
#[cfg(debug_assertions)]
fn evaluate_divisor<E: FieldElement>(
    divisor: &ConstraintDivisor<E::BaseField>,
    domain_size: usize,
    domain_offset: E::BaseField,
) -> Vec<E> {
    let g = E::BaseField::get_root_of_unity(domain_size.trailing_zeros());
    let domain = math::get_power_series_with_offset(g, domain_offset, domain_size);
    domain
        .into_iter()
        .map(|x| E::from(divisor.evaluate_at(x)))
        .collect()
}


// 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::{
    super::TraceLde, evaluation_table::EvaluationTableFragment, BoundaryConstraints,
    ConstraintEvaluationTable, PeriodicValueTable, StarkDomain,
};
use air::{
    Air, AuxTraceRandElements, ConstraintCompositionCoefficients, EvaluationFrame,
    TransitionConstraints,
};
use math::FieldElement;
use utils::iter_mut;

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

// CONSTANTS
// ================================================================================================

#[cfg(feature = "concurrent")]
const MIN_CONCURRENT_DOMAIN_SIZE: usize = 8192;

// CONSTRAINT EVALUATOR
// ================================================================================================

pub struct ConstraintEvaluator<'a, A: Air, E: FieldElement<BaseField = A::BaseField>> {
    air: &'a A,
    boundary_constraints: BoundaryConstraints<E>,
    transition_constraints: TransitionConstraints<E>,
    aux_rand_elements: AuxTraceRandElements<E>,
    periodic_values: PeriodicValueTable<E::BaseField>,
}

impl<'a, A: Air, E: FieldElement<BaseField = A::BaseField>> ConstraintEvaluator<'a, A, E> {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new evaluator which can be used to evaluate transition and boundary constraints
    /// over extended execution trace.
    pub fn new(
        air: &'a A,
        aux_rand_elements: AuxTraceRandElements<E>,
        composition_coefficients: ConstraintCompositionCoefficients<E>,
    ) -> Self {
        // build transition constraint groups; these will be used to compose transition constraint
        // evaluations
        let transition_constraints =
            air.get_transition_constraints(&composition_coefficients.transition);

        // build periodic value table
        let periodic_values = PeriodicValueTable::new(air);

        // build boundary constraint groups; these will be used to evaluate and compose boundary
        // constraint evaluations.
        let boundary_constraints =
            BoundaryConstraints::new(air, &aux_rand_elements, &composition_coefficients.boundary);

        ConstraintEvaluator {
            air,
            boundary_constraints,
            transition_constraints,
            aux_rand_elements,
            periodic_values,
        }
    }

    // EVALUATOR
    // --------------------------------------------------------------------------------------------
    /// Evaluates constraints against the provided extended execution trace. Constraints are
    /// evaluated over a constraint evaluation domain. This is an optimization because constraint
    /// evaluation domain can be many times smaller than the full LDE domain.
    pub fn evaluate(
        self,
        trace: &TraceLde<E>,
        domain: &'a StarkDomain<E::BaseField>,
    ) -> ConstraintEvaluationTable<'a, E> {
        assert_eq!(
            trace.trace_len(),
            domain.lde_domain_size(),
            "extended trace length is not consistent with evaluation domain"
        );

        // build a list of constraint divisors; currently, all transition constraints have the same
        // divisor which we put at the front of the list; boundary constraint divisors are appended
        // after that
        let mut divisors = vec![self.transition_constraints.divisor().clone()];
        divisors.append(&mut self.boundary_constraints.get_divisors());

        // allocate space for constraint evaluations; when we are in debug mode, we also allocate
        // memory to hold all transition constraint evaluations (before they are merged into a
        // single value) so that we can check their degrees later
        #[cfg(not(debug_assertions))]
        let mut evaluation_table = ConstraintEvaluationTable::<E>::new(domain, divisors);
        #[cfg(debug_assertions)]
        let mut evaluation_table =
            ConstraintEvaluationTable::<E>::new(domain, divisors, &self.transition_constraints);

        // when `concurrent` feature is enabled, break the evaluation table into multiple fragments
        // to evaluate them into multiple threads; unless the constraint evaluation domain is small,
        // then don't bother with concurrent evaluation

        #[cfg(not(feature = "concurrent"))]
        let num_fragments = 1;

        #[cfg(feature = "concurrent")]
        let num_fragments = if domain.ce_domain_size() >= MIN_CONCURRENT_DOMAIN_SIZE {
            rayon::current_num_threads().next_power_of_two()
        } else {
            1
        };

        // evaluate constraints for each fragment; if the trace consist of multiple segments
        // we evaluate constraints for all segments. otherwise, we evaluate constraints only
        // for the main segment.
        let mut fragments = evaluation_table.fragments(num_fragments);
        iter_mut!(fragments).for_each(|fragment| {
            if self.air.trace_info().is_multi_segment() {
                self.evaluate_fragment_full(trace, domain, fragment);
            } else {
                self.evaluate_fragment_main(trace, domain, fragment);
            }
        });

        // when in debug mode, make sure expected transition constraint degrees align with
        // actual degrees we got during constraint evaluation
        #[cfg(debug_assertions)]
        evaluation_table.validate_transition_degrees();

        evaluation_table
    }

    // EVALUATION HELPERS
    // --------------------------------------------------------------------------------------------

    /// Evaluates constraints for a single fragment of the evaluation table.
    ///
    /// This evaluates constraints only over the main segment of the execution trace.
    fn evaluate_fragment_main(
        &self,
        trace: &TraceLde<E>,
        domain: &StarkDomain<A::BaseField>,
        fragment: &mut EvaluationTableFragment<E>,
    ) {
        // initialize buffers to hold trace values and evaluation results at each step;
        let mut main_frame = EvaluationFrame::new(trace.main_trace_width());
        let mut evaluations = vec![E::ZERO; fragment.num_columns()];
        let mut t_evaluations = vec![E::BaseField::ZERO; self.num_main_transition_constraints()];

        // this will be used to convert steps in constraint evaluation domain to steps in
        // LDE domain
        let lde_shift = domain.ce_to_lde_blowup().trailing_zeros();

        for i in 0..fragment.num_rows() {
            let step = i + fragment.offset();

            // update evaluation frame buffer with data from the execution trace; this will
            // read current and next rows from the trace into the buffer; data in the trace
            // table is extended over the LDE domain, so, we need to convert step in constraint
            // evaluation domain, into a step in LDE domain, in case these domains are different
            trace.read_main_trace_frame_into(step << lde_shift, &mut main_frame);

            // evaluate transition constraints and save the merged result the first slot of the
            // evaluations buffer
            evaluations[0] = self.evaluate_main_transition(&main_frame, step, &mut t_evaluations);

            // when in debug mode, save transition constraint evaluations
            #[cfg(debug_assertions)]
            fragment.update_transition_evaluations(i, &t_evaluations, &[]);

            // evaluate boundary constraints; the results go into remaining slots of the
            // evaluations buffer
            let main_state = main_frame.current();
            self.boundary_constraints.evaluate_main(
                main_state,
                domain,
                step,
                &mut evaluations[1..],
            );

            // record the result in the evaluation table
            fragment.update_row(i, &evaluations);
        }
    }

    /// Evaluates constraints for a single fragment of the evaluation table.
    ///
    /// This evaluates constraints only over all segments of the execution trace (i.e. main segment
    /// and all auxiliary segments).
    fn evaluate_fragment_full(
        &self,
        trace: &TraceLde<E>,
        domain: &StarkDomain<A::BaseField>,
        fragment: &mut EvaluationTableFragment<E>,
    ) {
        // initialize buffers to hold trace values and evaluation results at each step
        let mut main_frame = EvaluationFrame::new(trace.main_trace_width());
        let mut aux_frame = EvaluationFrame::new(trace.aux_trace_width());
        let mut tm_evaluations = vec![E::BaseField::ZERO; self.num_main_transition_constraints()];
        let mut ta_evaluations = vec![E::ZERO; self.num_aux_transition_constraints()];
        let mut evaluations = vec![E::ZERO; fragment.num_columns()];

        // this will be used to convert steps in constraint evaluation domain to steps in
        // LDE domain
        let lde_shift = domain.ce_to_lde_blowup().trailing_zeros();

        for i in 0..fragment.num_rows() {
            let step = i + fragment.offset();

            // read both the main and the auxiliary evaluation frames from the trace
            trace.read_main_trace_frame_into(step << lde_shift, &mut main_frame);
            trace.read_aux_trace_frame_into(step << lde_shift, &mut aux_frame);

            // evaluate transition constraints and save the merged result the first slot of the
            // evaluations buffer; we evaluate and compose constraints in the same function, we
            // can just add up the results of evaluating main and auxiliary constraints.
            evaluations[0] = self.evaluate_main_transition(&main_frame, step, &mut tm_evaluations);
            evaluations[0] +=
                self.evaluate_aux_transition(&main_frame, &aux_frame, step, &mut ta_evaluations);

            // when in debug mode, save transition constraint evaluations
            #[cfg(debug_assertions)]
            fragment.update_transition_evaluations(i, &tm_evaluations, &ta_evaluations);

            // evaluate boundary constraints; the results go into remaining slots of the
            // evaluations buffer
            let main_state = main_frame.current();
            let aux_state = aux_frame.current();
            self.boundary_constraints.evaluate_all(
                main_state,
                aux_state,
                domain,
                step,
                &mut evaluations[1..],
            );

            // record the result in the evaluation table
            fragment.update_row(i, &evaluations);
        }
    }

    // TRANSITION CONSTRAINT EVALUATORS
    // --------------------------------------------------------------------------------------------

    /// Evaluates transition constraints of the main execution trace at the specified step of the
    /// constraint evaluation domain.
    ///
    /// `x` is the corresponding domain value at the specified step. That is, x = s * g^step,
    /// where g is the generator of the constraint evaluation domain, and s is the domain offset.
    #[rustfmt::skip]
    fn evaluate_main_transition(
        &self,
        main_frame: &EvaluationFrame<E::BaseField>,
        step: usize,
        evaluations: &mut [E::BaseField],
    ) -> E {
        // TODO: use a more efficient way to zero out memory
        evaluations.fill(E::BaseField::ZERO);

        // get periodic values at the evaluation step
        let periodic_values = self.periodic_values.get_row(step);

        // evaluate transition constraints over the main segment of the execution trace and save
        // the results into evaluations buffer
        self.air.evaluate_transition(main_frame, periodic_values, evaluations);

        // merge transition constraint evaluations into a single value and return it;
        // we can do this here because all transition constraints have the same divisor.
        evaluations
            .iter()
            .zip(self.transition_constraints.main_constraint_coef().iter())
            .fold(E::ZERO, |acc, (&const_eval, &coef)| {
                acc + coef.mul_base(const_eval)
            })
    }

    /// Evaluates all transition constraints (i.e., for main and auxiliary trace segments) at the
    /// specified step of the constraint evaluation domain.
    ///
    /// `x` is the corresponding domain value at the specified step. That is, x = s * g^step,
    /// where g is the generator of the constraint evaluation domain, and s is the domain offset.
    #[rustfmt::skip]
    fn evaluate_aux_transition(
        &self,
        main_frame: &EvaluationFrame<E::BaseField>,
        aux_frame: &EvaluationFrame<E>,
        step: usize,
        evaluations: &mut [E],
    ) -> E {
        // TODO: use a more efficient way to zero out memory
        evaluations.fill(E::ZERO);

        // get periodic values at the evaluation step
        let periodic_values = self.periodic_values.get_row(step);

        // evaluate transition constraints over auxiliary trace segments and save the results into
        // evaluations buffer
        self.air.evaluate_aux_transition(
            main_frame,
            aux_frame,
            periodic_values,
            &self.aux_rand_elements,
            evaluations,
        );

        // merge transition constraint evaluations into a single value and return it;
        // we can do this here because all transition constraints have the same divisor.
        evaluations
            .iter()
            .zip(self.transition_constraints.aux_constraint_coef().iter())
            .fold(E::ZERO, |acc, (&const_eval, &coef)| acc + coef * const_eval)
    }

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

    /// Returns the number of transition constraints applied against the main segment of the
    /// execution trace.
    fn num_main_transition_constraints(&self) -> usize {
        self.transition_constraints.num_main_constraints()
    }

    /// Returns the number of transition constraints applied against all auxiliary trace segments.
    fn num_aux_transition_constraints(&self) -> usize {
        self.transition_constraints.num_aux_constraints()
    }
}


// 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::{ColMatrix, ConstraintDivisor, ProverError, RowMatrix, StarkDomain};

mod boundary;
use boundary::BoundaryConstraints;

mod periodic_table;
use periodic_table::PeriodicValueTable;

mod evaluator;
pub use evaluator::ConstraintEvaluator;

mod composition_poly;
pub use composition_poly::CompositionPoly;

mod evaluation_table;
pub use evaluation_table::ConstraintEvaluationTable;

mod commitment;
pub use commitment::ConstraintCommitment;


// 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 air::Air;
use math::{fft, StarkField};
use utils::{
    collections::{BTreeMap, Vec},
    uninit_vector,
};

pub struct PeriodicValueTable<B: StarkField> {
    values: Vec<B>,
    length: usize,
    width: usize,
}

impl<B: StarkField> PeriodicValueTable<B> {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Builds a table of periodic column values for the specified AIR. The table contains expanded
    /// values of all periodic columns normalized to the same length. This enables simple lookup
    /// into the able using step index of the constraint evaluation domain.
    pub fn new<A: Air<BaseField = B>>(air: &A) -> PeriodicValueTable<B> {
        // get a list of polynomials describing periodic columns from AIR. if there are no
        // periodic columns return an empty table
        let polys = air.get_periodic_column_polys();
        if polys.is_empty() {
            return PeriodicValueTable {
                values: Vec::new(),
                length: 0,
                width: 0,
            };
        }

        // determine the size of the biggest polynomial in the set. unwrap is OK here
        // because if we get here, there must be at least one polynomial in the set.
        let max_poly_size = polys.iter().max_by_key(|p| p.len()).unwrap().len();

        // cache twiddles used for polynomial evaluation here so that we don't have to re-build
        // them for polynomials of the same size
        let mut twiddle_map = BTreeMap::new();

        let evaluations = polys
            .iter()
            .map(|poly| {
                let poly_size = poly.len();
                let num_cycles = (air.trace_length() / poly_size) as u64;
                let offset = air.domain_offset().exp(num_cycles.into());
                let twiddles = twiddle_map
                    .entry(poly_size)
                    .or_insert_with(|| fft::get_twiddles(poly_size));

                fft::evaluate_poly_with_offset(poly, twiddles, offset, air.ce_blowup_factor())
            })
            .collect::<Vec<_>>();

        // allocate memory to hold all expanded values and copy polynomial evaluations into the
        // table in such a way that values for the same row are adjacent to each other.
        let row_width = polys.len();
        let column_length = max_poly_size * air.ce_blowup_factor();
        let mut values = unsafe { uninit_vector(row_width * column_length) };
        for i in 0..column_length {
            for (j, column) in evaluations.iter().enumerate() {
                values[i * row_width + j] = column[i % column.len()];
            }
        }

        PeriodicValueTable {
            values,
            length: column_length,
            width: row_width,
        }
    }

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

    pub fn is_empty(&self) -> bool {
        self.width == 0
    }

    pub fn get_row(&self, ce_step: usize) -> &[B] {
        if self.is_empty() {
            &[]
        } else {
            let start = (ce_step % self.length) * self.width;
            &self.values[start..start + self.width]
        }
    }
}

// TESTS
// ================================================================================================

#[cfg(test)]
mod tests {
    use crate::tests::MockAir;
    use air::Air;
    use math::{
        fields::f128::BaseElement, get_power_series_with_offset, polynom, FieldElement, StarkField,
    };
    use utils::collections::Vec;

    #[test]
    fn periodic_value_table() {
        let trace_length = 32;

        // instantiate AIR with 2 periodic columns
        let col1 = vec![1u128, 2]
            .into_iter()
            .map(BaseElement::new)
            .collect::<Vec<_>>();
        let col2 = vec![3u128, 4, 5, 6]
            .into_iter()
            .map(BaseElement::new)
            .collect::<Vec<_>>();
        let air = MockAir::with_periodic_columns(vec![col1, col2], trace_length);

        // build a table of periodic values
        let table = super::PeriodicValueTable::new(&air);

        assert_eq!(2, table.width);
        assert_eq!(4 * air.ce_blowup_factor(), table.length);

        let polys = air.get_periodic_column_polys();
        let domain = build_ce_domain(air.ce_domain_size(), air.domain_offset());

        // build expected values by evaluating polynomials over shifted ce_domain
        let expected = polys
            .iter()
            .map(|poly| {
                let num_cycles = trace_length / poly.len();
                domain
                    .iter()
                    .map(|&x| {
                        let x = x.exp((num_cycles as u32).into());
                        polynom::eval(poly, x)
                    })
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>();

        // build actual values by recording rows of the table at each step of ce_domain
        let mut actual = vec![Vec::new(), Vec::new()];
        for i in 0..air.ce_domain_size() {
            let row = table.get_row(i);
            actual[0].push(row[0]);
            actual[1].push(row[1]);
        }

        assert_eq!(expected, actual);
    }

    fn build_ce_domain(domain_size: usize, domain_offset: BaseElement) -> Vec<BaseElement> {
        let g = BaseElement::get_root_of_unity(domain_size.ilog2());
        get_power_series_with_offset(g, domain_offset, 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 crate::StarkDomain;
use core::{iter::FusedIterator, slice};
use crypto::{ElementHasher, MerkleTree};
use math::{fft, polynom, FieldElement};
use utils::{batch_iter_mut, collections::Vec, iter, iter_mut, uninit_vector};

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

// COLUMN-MAJOR MATRIX
// ================================================================================================

/// A two-dimensional matrix of field elements arranged in column-major order.
///
/// This struct is used as a backing type for many objects manipulated by the prover. The matrix
/// itself does not assign any contextual meaning to the values stored in it. For example, columns
/// may contain evaluations of polynomials, or polynomial coefficients, or really anything else.
/// However, the matrix does expose a number of methods which make assumptions about the underlying
/// data.
///
/// A matrix imposes the following restrictions on its content:
/// - A matrix must consist of at least 1 column and at least 2 rows.
/// - All columns must be of the same length.
/// - Number of rows must be a power of two.
#[derive(Debug, Clone)]
pub struct ColMatrix<E: FieldElement> {
    columns: Vec<Vec<E>>,
}

impl<E: FieldElement> ColMatrix<E> {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new [Matrix] instantiated with the data from the specified columns.
    ///
    /// # Panics
    /// Panics if:
    /// * The provided vector of columns is empty.
    /// * Not all of the columns have the same number of elements.
    /// * Number of rows is smaller than or equal to 1.
    /// * Number of rows is not a power of two.
    pub fn new(columns: Vec<Vec<E>>) -> Self {
        assert!(
            !columns.is_empty(),
            "a matrix must contain at least one column"
        );
        let num_rows = columns[0].len();
        assert!(
            num_rows > 1,
            "number of rows in a matrix must be greater than one"
        );
        assert!(
            num_rows.is_power_of_two(),
            "number of rows in a matrix must be a power of 2"
        );
        for column in columns.iter().skip(1) {
            assert_eq!(
                column.len(),
                num_rows,
                "all matrix columns must have the same length"
            );
        }

        Self { columns }
    }

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

    /// Returns the number of columns in this matrix.
    pub fn num_cols(&self) -> usize {
        self.columns.len()
    }

    /// Returns the number of base field columns in this matrix.
    ///
    /// The number of base field columns is defined as the number of columns multiplied by the
    /// extension degree of field elements contained in this matrix.
    pub fn num_base_cols(&self) -> usize {
        self.num_cols() * E::EXTENSION_DEGREE
    }

    /// Returns the number of rows in this matrix.
    pub fn num_rows(&self) -> usize {
        self.columns[0].len()
    }

    /// Returns the element located at the specified column and row indexes in this matrix.
    ///
    /// # Panics
    /// Panics if either `col_idx` or `row_idx` are out of bounds for this matrix.
    pub fn get(&self, col_idx: usize, row_idx: usize) -> E {
        self.columns[col_idx][row_idx]
    }

    /// Returns base field elements located at the specified column and row indexes in this matrix.
    ///
    /// For STARK fields, `base_col_idx` is the same as `col_idx` used in `Self::get` method. For
    /// extension fields, each column in the matrix is viewed as 2 or more columns in the base
    /// field.
    ///
    /// Thus, for example, if we are in a degree 2 extension field, `base_col_idx = 0` would refer
    /// to the first base element of the first column, `base_col_idx = 1` would refer to the second
    /// base element of the first column, `base_col_idx = 2` would refer to the first base element
    /// of the second column etc.
    ///
    /// # Panics
    /// Panics if either `base_col_idx` or `row_idx` are out of bounds for this matrix.
    pub fn get_base_element(&self, base_col_idx: usize, row_idx: usize) -> E::BaseField {
        let (col_idx, elem_idx) = (
            base_col_idx / E::EXTENSION_DEGREE,
            base_col_idx % E::EXTENSION_DEGREE,
        );
        self.columns[col_idx][row_idx].base_element(elem_idx)
    }

    /// Set the cell in this matrix at the specified column and row indexes to the provided value.
    ///
    /// # Panics
    /// Panics if either `col_idx` or `row_idx` are out of bounds for this matrix.
    pub fn set(&mut self, col_idx: usize, row_idx: usize, value: E) {
        self.columns[col_idx][row_idx] = value;
    }

    /// Returns a reference to the column at the specified index.
    pub fn get_column(&self, col_idx: usize) -> &[E] {
        &self.columns[col_idx]
    }

    /// Returns a reference to the column at the specified index.
    pub fn get_column_mut(&mut self, col_idx: usize) -> &mut [E] {
        &mut self.columns[col_idx]
    }

    /// Copies values of all columns at the specified row into the specified row slice.
    ///
    /// # Panics
    /// Panics if `row_idx` is out of bounds for this matrix.
    pub fn read_row_into(&self, row_idx: usize, row: &mut [E]) {
        for (column, value) in self.columns.iter().zip(row.iter_mut()) {
            *value = column[row_idx];
        }
    }

    /// Updates a row in this matrix at the specified index to the provided data.
    ///
    /// # Panics
    /// Panics if `row_idx` is out of bounds for this matrix.
    pub fn update_row(&mut self, row_idx: usize, row: &[E]) {
        for (column, &value) in self.columns.iter_mut().zip(row) {
            column[row_idx] = value;
        }
    }

    /// Merges a column to the end of the matrix provided its length matches the matrix.
    ///
    /// # Panics
    /// Panics if the column has a different length to other columns in the matrix.
    pub fn merge_column(&mut self, column: Vec<E>) {
        if let Some(first_column) = self.columns.first() {
            assert_eq!(first_column.len(), column.len());
        }
        self.columns.push(column);
    }

    // ITERATION
    // --------------------------------------------------------------------------------------------

    /// Returns an iterator over the columns of this matrix.
    pub fn columns(&self) -> ColumnIter<E> {
        ColumnIter::new(self)
    }

    /// Returns a mutable iterator over the columns of this matrix.
    pub fn columns_mut(&mut self) -> ColumnIterMut<E> {
        ColumnIterMut::new(self)
    }

    // POLYNOMIAL METHODS
    // --------------------------------------------------------------------------------------------

    /// Interpolates columns of the matrix into polynomials in coefficient form and returns the
    /// result.
    ///
    /// The interpolation is performed as follows:
    /// * Each column of the matrix is interpreted as evaluations of degree `num_rows - 1`
    ///   polynomial over a subgroup of size `num_rows`.
    /// * Then each column is interpolated using iFFT algorithm into a polynomial in coefficient
    ///   form.
    /// * The resulting polynomials are returned as a single matrix where each column contains
    ///   coefficients of a degree `num_rows - 1` polynomial.
    pub fn interpolate_columns(&self) -> Self {
        let inv_twiddles = fft::get_inv_twiddles::<E::BaseField>(self.num_rows());
        let columns = iter!(self.columns)
            .map(|evaluations| {
                let mut column = evaluations.clone();
                fft::interpolate_poly(&mut column, &inv_twiddles);
                column
            })
            .collect();
        Self { columns }
    }

    /// Interpolates columns of the matrix into polynomials in coefficient form and returns the
    /// result. The input matrix is consumed in the process.
    ///
    /// The interpolation is performed as follows:
    /// * Each column of the matrix is interpreted as evaluations of degree `num_rows - 1`
    ///   polynomial over a subgroup of size `num_rows`.
    /// * Then each column is interpolated (in place) using iFFT algorithm into a polynomial in
    ///   coefficient form.
    /// * The resulting polynomials are returned as a single matrix where each column contains
    ///   coefficients of a degree `num_rows - 1` polynomial.
    pub fn interpolate_columns_into(mut self) -> Self {
        let inv_twiddles = fft::get_inv_twiddles::<E::BaseField>(self.num_rows());
        iter_mut!(self.columns).for_each(|column| fft::interpolate_poly(column, &inv_twiddles));
        self
    }

    /// Evaluates polynomials contained in the columns of this matrix over the specified domain
    /// and returns the result.
    ///
    /// The evaluation is done as follows:
    /// * Each column of the matrix is interpreted as coefficients of degree `num_rows - 1`
    ///   polynomial.
    /// * These polynomials are evaluated over the LDE domain defined by the specified
    ///   [StarkDomain] using FFT algorithm. The domain specification includes the size of the
    ///   subgroup as well as the domain offset (to define a coset).
    /// * The resulting evaluations are returned in a new Matrix.
    pub fn evaluate_columns_over(&self, domain: &StarkDomain<E::BaseField>) -> Self {
        let columns = iter!(self.columns)
            .map(|poly| {
                fft::evaluate_poly_with_offset(
                    poly,
                    domain.trace_twiddles(),
                    domain.offset(),
                    domain.trace_to_lde_blowup(),
                )
            })
            .collect();
        Self { columns }
    }

    /// Evaluates polynomials contained in the columns of this matrix at a single point `x`.
    pub fn evaluate_columns_at<F>(&self, x: F) -> Vec<F>
    where
        F: FieldElement + From<E>,
    {
        iter!(self.columns).map(|p| polynom::eval(p, x)).collect()
    }

    // COMMITMENTS
    // --------------------------------------------------------------------------------------------

    /// Returns a commitment to this matrix.
    ///
    /// The commitment is built as follows:
    /// * Each row of the matrix is hashed into a single digest of the specified hash function.
    /// * The resulting values are used to built a binary Merkle tree such that each row digest
    ///   becomes a leaf in the tree. Thus, the number of leaves in the tree is equal to the
    ///   number of rows in the matrix.
    /// * The resulting Merkle tree is return as the commitment to the entire matrix.
    pub fn commit_to_rows<H>(&self) -> MerkleTree<H>
    where
        H: ElementHasher<BaseField = E::BaseField>,
    {
        // allocate vector to store row hashes
        let mut row_hashes = unsafe { uninit_vector::<H::Digest>(self.num_rows()) };

        // iterate though matrix rows, hashing each row; the hashing is done by first copying a
        // row into row_buf to avoid heap allocations, and then by applying the hash function to
        // the buffer.
        batch_iter_mut!(
            &mut row_hashes,
            128, // min batch size
            |batch: &mut [H::Digest], batch_offset: usize| {
                let mut row_buf = vec![E::ZERO; self.num_cols()];
                for (i, row_hash) in batch.iter_mut().enumerate() {
                    self.read_row_into(i + batch_offset, &mut row_buf);
                    *row_hash = H::hash_elements(&row_buf);
                }
            }
        );

        // build Merkle tree out of hashed rows
        MerkleTree::new(row_hashes).expect("failed to construct trace Merkle tree")
    }

    // CONVERSIONS
    // --------------------------------------------------------------------------------------------

    /// Returns the columns of this matrix as a list of vectors.
    ///
    /// TODO: replace this with an iterator.
    pub fn into_columns(self) -> Vec<Vec<E>> {
        self.columns
    }
}

// COLUMN ITERATOR
// ================================================================================================

pub struct ColumnIter<'a, E: FieldElement> {
    matrix: &'a ColMatrix<E>,
    cursor: usize,
}

impl<'a, E: FieldElement> ColumnIter<'a, E> {
    pub fn new(matrix: &'a ColMatrix<E>) -> Self {
        Self { matrix, cursor: 0 }
    }
}

impl<'a, E: FieldElement> Iterator for ColumnIter<'a, E> {
    type Item = &'a [E];

    fn next(&mut self) -> Option<Self::Item> {
        match self.matrix.num_cols() - self.cursor {
            0 => None,
            _ => {
                let column = self.matrix.get_column(self.cursor);
                self.cursor += 1;
                Some(column)
            }
        }
    }
}

impl<'a, E: FieldElement> ExactSizeIterator for ColumnIter<'a, E> {
    fn len(&self) -> usize {
        self.matrix.num_cols()
    }
}

impl<'a, E: FieldElement> FusedIterator for ColumnIter<'a, E> {}

// MUTABLE COLUMN ITERATOR
// ================================================================================================

pub struct ColumnIterMut<'a, E: FieldElement> {
    matrix: &'a mut ColMatrix<E>,
    cursor: usize,
}

impl<'a, E: FieldElement> ColumnIterMut<'a, E> {
    pub fn new(matrix: &'a mut ColMatrix<E>) -> Self {
        Self { matrix, cursor: 0 }
    }
}

impl<'a, E: FieldElement> Iterator for ColumnIterMut<'a, E> {
    type Item = &'a mut [E];

    fn next(&mut self) -> Option<Self::Item> {
        match self.matrix.num_cols() - self.cursor {
            0 => None,
            _ => {
                let column = self.matrix.get_column_mut(self.cursor);
                self.cursor += 1;

                // this is needed to get around mutable iterator lifetime issues; this is safe
                // because the iterator can never yield a reference to the same column twice
                let p = column.as_ptr();
                let len = column.len();
                Some(unsafe { slice::from_raw_parts_mut(p as *mut E, len) })
            }
        }
    }
}

impl<'a, E: FieldElement> ExactSizeIterator for ColumnIterMut<'a, E> {
    fn len(&self) -> usize {
        self.matrix.num_cols()
    }
}

impl<'a, E: FieldElement> FusedIterator for ColumnIterMut<'a, E> {}

// MULTI-MATRIX COLUMN ITERATOR
// ================================================================================================

pub struct MultiColumnIter<'a, E: FieldElement> {
    matrixes: &'a [ColMatrix<E>],
    m_cursor: usize,
    c_cursor: usize,
}

impl<'a, E: FieldElement> MultiColumnIter<'a, E> {
    pub fn new(matrixes: &'a [ColMatrix<E>]) -> Self {
        // make sure all matrixes have the same number of rows
        if !matrixes.is_empty() {
            let num_rows = matrixes[0].num_rows();
            for matrix in matrixes.iter().skip(1) {
                assert_eq!(
                    matrix.num_rows(),
                    num_rows,
                    "all matrixes must have the same number of rows"
                );
            }
        }

        Self {
            matrixes,
            m_cursor: 0,
            c_cursor: 0,
        }
    }
}

impl<'a, E: FieldElement> Iterator for MultiColumnIter<'a, E> {
    type Item = &'a [E];

    fn next(&mut self) -> Option<Self::Item> {
        if self.matrixes.is_empty() {
            return None;
        }
        let matrix = &self.matrixes[self.m_cursor];
        match matrix.num_cols() - self.c_cursor {
            0 => None,
            _ => {
                let column = matrix.get_column(self.c_cursor);
                self.c_cursor += 1;
                if self.c_cursor == matrix.num_cols() && self.m_cursor < self.matrixes.len() - 1 {
                    self.m_cursor += 1;
                    self.c_cursor = 0;
                }
                Some(column)
            }
        }
    }
}

impl<'a, E: FieldElement> ExactSizeIterator for MultiColumnIter<'a, E> {
    fn len(&self) -> usize {
        self.matrixes.iter().fold(0, |s, m| s + m.num_cols())
    }
}

impl<'a, E: FieldElement> FusedIterator for MultiColumnIter<'a, E> {}


// 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.

mod row_matrix;
pub use row_matrix::{build_segments, get_evaluation_offsets, RowMatrix};

mod col_matrix;
pub use col_matrix::{ColMatrix, ColumnIter, MultiColumnIter};

mod segments;
pub use segments::Segment;

#[cfg(test)]
mod tests;


// 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::{ColMatrix, Segment};
use crate::StarkDomain;
use crypto::{ElementHasher, MerkleTree};
use math::{fft, FieldElement, StarkField};
use utils::collections::Vec;
use utils::{batch_iter_mut, flatten_vector_elements, uninit_vector};

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

// ROW-MAJOR MATRIX
// ================================================================================================

/// A two-dimensional matrix of field elements arranged in row-major order.
///
/// The matrix is represented as a single vector of base field elements for the field defined by E
/// type parameter. The first `row_width` base field elements represent the first row of the matrix,
/// the next `row_width` base field elements represent the second row, and so on.
///
/// When rows are returned via the [RowMatrix::row()] method, base field elements are grouped
/// together as appropriate to form elements in E.
///
/// In some cases, rows may be padded with extra elements. The number of elements which are
/// accessible via the [RowMatrix::row()] method is specified by the `elements_per_row` member.
#[derive(Clone, Debug)]
pub struct RowMatrix<E: FieldElement> {
    /// Field elements stored in the matrix.
    data: Vec<E::BaseField>,
    /// Total number of base field elements stored in a single row.
    row_width: usize,
    /// Number of field elements in a single row accessible via the [RowMatrix::row()] method. This
    /// must be equal to or smaller than `row_width`.
    elements_per_row: usize,
}

impl<E: FieldElement> RowMatrix<E> {
    // CONSTRUCTORS
    // --------------------------------------------------------------------------------------------

    /// Returns a new [RowMatrix] constructed by evaluating the provided polynomials over the
    /// domain defined by the specified blowup factor.
    ///
    /// The provided `polys` matrix is assumed to contain polynomials in coefficient form (one
    /// polynomial per column). Columns in the returned matrix will contain evaluations of the
    /// corresponding polynomials over the domain defined by polynomial size (i.e., number of rows
    /// in the `polys` matrix) and the `blowup_factor`.
    ///
    /// To improve performance, polynomials are evaluated in batches specified by the `N` type
    /// parameter. Minimum batch size is 1.
    pub fn evaluate_polys<const N: usize>(polys: &ColMatrix<E>, blowup_factor: usize) -> Self {
        assert!(N > 0, "batch size N must be greater than zero");

        // pre-compute offsets for each row
        let poly_size = polys.num_rows();
        let offsets =
            get_evaluation_offsets::<E>(poly_size, blowup_factor, E::BaseField::GENERATOR);

        // compute twiddles for polynomial evaluation
        let twiddles = fft::get_twiddles::<E::BaseField>(polys.num_rows());

        // build matrix segments by evaluating all polynomials
        let segments = build_segments::<E, N>(polys, &twiddles, &offsets);

        // transpose data in individual segments into a single row-major matrix
        Self::from_segments(segments, polys.num_base_cols())
    }

    /// Returns a new [RowMatrix] constructed by evaluating the provided polynomials over the
    /// specified [StarkDomain].
    ///
    /// The provided `polys` matrix is assumed to contain polynomials in coefficient form (one
    /// polynomial per column). Columns in the returned matrix will contain evaluations of the
    /// corresponding polynomials over the LDE domain defined by the provided [StarkDomain].
    ///
    /// To improve performance, polynomials are evaluated in batches specified by the `N` type
    /// parameter. Minimum batch size is 1.
    pub fn evaluate_polys_over<const N: usize>(
        polys: &ColMatrix<E>,
        domain: &StarkDomain<E::BaseField>,
    ) -> Self {
        assert!(N > 0, "batch size N must be greater than zero");

        // pre-compute offsets for each row
        let poly_size = polys.num_rows();
        let offsets =
            get_evaluation_offsets::<E>(poly_size, domain.trace_to_lde_blowup(), domain.offset());

        // build matrix segments by evaluating all polynomials
        let segments = build_segments::<E, N>(polys, domain.trace_twiddles(), &offsets);

        // transpose data in individual segments into a single row-major matrix
        Self::from_segments(segments, polys.num_base_cols())
    }

    /// Returns a new [RowMatrix] instantiated from the specified matrix segments.
    ///
    /// `elements_per_row` specifies how many base field elements are considered to form a single
    /// row in the matrix.
    ///
    /// # Panics
    /// Panics if
    /// - `segments` is an empty vector.
    /// - `elements_per_row` is greater than the row width implied by the number of segments and
    ///   `N` type parameter.
    pub fn from_segments<const N: usize>(
        segments: Vec<Segment<E::BaseField, N>>,
        elements_per_row: usize,
    ) -> Self {
        assert!(N > 0, "batch size N must be greater than zero");
        assert!(!segments.is_empty(), "a list of segments cannot be empty");

        // compute the size of each row
        let row_width = segments.len() * N;
        assert!(
            elements_per_row <= row_width,
            "elements per row cannot exceed {row_width}, but was {elements_per_row}"
        );

        // transpose the segments into a single vector of arrays
        let result = transpose(segments);

        // flatten the result to be a simple vector of elements and return
        RowMatrix {
            data: flatten_vector_elements(result),
            row_width,
            elements_per_row,
        }
    }

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

    /// Returns the number of columns in this matrix.
    pub fn num_cols(&self) -> usize {
        self.elements_per_row / E::EXTENSION_DEGREE
    }

    /// Returns the number of rows in this matrix.
    pub fn num_rows(&self) -> usize {
        self.data.len() / self.row_width
    }

    /// Returns the element located at the specified column and row indexes in this matrix.
    ///
    /// # Panics
    /// Panics if either `col_idx` or `row_idx` are out of bounds for this matrix.
    pub fn get(&self, col_idx: usize, row_idx: usize) -> E {
        self.row(row_idx)[col_idx]
    }

    /// Returns a reference to a row at the specified index in this matrix.
    ///
    /// # Panics
    /// Panics if the specified row index is out of bounds.
    pub fn row(&self, row_idx: usize) -> &[E] {
        assert!(row_idx < self.num_rows());
        let start = row_idx * self.row_width;
        E::slice_from_base_elements(&self.data[start..start + self.elements_per_row])
    }

    /// Returns the data in this matrix as a slice of field elements.
    pub fn data(&self) -> &[E::BaseField] {
        &self.data
    }

    // COMMITMENTS
    // --------------------------------------------------------------------------------------------

    /// Returns a commitment to this matrix.
    ///
    /// The commitment is built as follows:
    /// * Each row of the matrix is hashed into a single digest of the specified hash function.
    /// * The resulting values are used to build a binary Merkle tree such that each row digest
    ///   becomes a leaf in the tree. Thus, the number of leaves in the tree is equal to the
    ///   number of rows in the matrix.
    /// * The resulting Merkle tree is returned as the commitment to the entire matrix.
    pub fn commit_to_rows<H>(&self) -> MerkleTree<H>
    where
        H: ElementHasher<BaseField = E::BaseField>,
    {
        // allocate vector to store row hashes
        let mut row_hashes = unsafe { uninit_vector::<H::Digest>(self.num_rows()) };

        // iterate though matrix rows, hashing each row
        batch_iter_mut!(
            &mut row_hashes,
            128, // min batch size
            |batch: &mut [H::Digest], batch_offset: usize| {
                for (i, row_hash) in batch.iter_mut().enumerate() {
                    *row_hash = H::hash_elements(self.row(batch_offset + i));
                }
            }
        );

        // build Merkle tree out of hashed rows
        MerkleTree::new(row_hashes).expect("failed to construct trace Merkle tree")
    }
}

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

/// Returns a vector of offsets for an evaluation defined by the specified polynomial size, blowup
/// factor and domain offset.
///
/// When `concurrent` feature is enabled, offsets are computed in multiple threads.
pub fn get_evaluation_offsets<E: FieldElement>(
    poly_size: usize,
    blowup_factor: usize,
    domain_offset: E::BaseField,
) -> Vec<E::BaseField> {
    let domain_size = poly_size * blowup_factor;
    let g = E::BaseField::get_root_of_unity(domain_size.ilog2());

    // allocate memory to hold the offsets
    let mut offsets = unsafe { uninit_vector(domain_size) };

    // define a closure to compute offsets for a given chunk of the result; the number of chunks
    // is defined by the blowup factor. for example, for blowup factor = 2, the number of chunks
    // will be 2, for blowup factor = 8, the number of chunks will be 8 etc.
    let compute_offsets = |(chunk_idx, chunk): (usize, &mut [E::BaseField])| {
        let idx = fft::permute_index(blowup_factor, chunk_idx) as u64;
        let offset = g.exp_vartime(idx.into()) * domain_offset;
        let mut factor = E::BaseField::ONE;
        for res in chunk.iter_mut() {
            *res = factor;
            factor *= offset;
        }
    };

    // compute offsets for each chunk using either parallel or regular iterators

    #[cfg(not(feature = "concurrent"))]
    offsets
        .chunks_mut(poly_size)
        .enumerate()
        .for_each(compute_offsets);

    #[cfg(feature = "concurrent")]
    offsets
        .par_chunks_mut(poly_size)
        .enumerate()
        .for_each(compute_offsets);

    offsets
}

/// Returns matrix segments constructed by evaluating polynomials in the specified matrix over the
/// domain defined by twiddles and offsets.
pub fn build_segments<E: FieldElement, const N: usize>(
    polys: &ColMatrix<E>,
    twiddles: &[E::BaseField],
    offsets: &[E::BaseField],
) -> Vec<Segment<E::BaseField, N>> {
    assert!(N > 0, "batch size N must be greater than zero");
    debug_assert_eq!(polys.num_rows(), twiddles.len() * 2);
    debug_assert_eq!(offsets.len() % polys.num_rows(), 0);

    let num_segments = if polys.num_base_cols() % N == 0 {
        polys.num_base_cols() / N
    } else {
        polys.num_base_cols() / N + 1
    };

    (0..num_segments)
        .map(|i| Segment::new(polys, i * N, offsets, twiddles))
        .collect()
}

/// Transposes a vector of segments into a single vector of fixed-size arrays.
///
/// When `concurrent` feature is enabled, transposition is performed in multiple threads.
fn transpose<B: StarkField, const N: usize>(mut segments: Vec<Segment<B, N>>) -> Vec<[B; N]> {
    let num_rows = segments[0].num_rows();
    let num_segs = segments.len();
    let result_len = num_rows * num_segs;

    // if there is only one segment, there is nothing to transpose as it is already in row
    // major form
    if segments.len() == 1 {
        return segments.remove(0).into_data();
    }

    // allocate memory to hold the transposed result;
    // TODO: investigate transposing in-place
    let mut result = unsafe { uninit_vector::<[B; N]>(result_len) };

    // determine number of batches in which transposition will be preformed; if `concurrent`
    // feature is not enabled, the number of batches will always be 1
    let num_batches = get_num_batches(result_len);
    let rows_per_batch = num_rows / num_batches;

    // define a closure for transposing a given batch
    let transpose_batch = |(batch_idx, batch): (usize, &mut [[B; N]])| {
        let row_offset = batch_idx * rows_per_batch;
        for i in 0..rows_per_batch {
            let row_idx = i + row_offset;
            for j in 0..num_segs {
                let v = &segments[j][row_idx];
                batch[i * num_segs + j].copy_from_slice(v);
            }
        }
    };

    // call the closure either once (for single-threaded transposition) or in a parallel
    // iterator (for multi-threaded transposition)

    #[cfg(not(feature = "concurrent"))]
    transpose_batch((0, &mut result));

    #[cfg(feature = "concurrent")]
    result
        .par_chunks_mut(result_len / num_batches)
        .enumerate()
        .for_each(transpose_batch);

    result
}

#[cfg(not(feature = "concurrent"))]
fn get_num_batches(_input_size: usize) -> usize {
    1
}

#[cfg(feature = "concurrent")]
fn get_num_batches(input_size: usize) -> usize {
    if input_size < 1024 {
        return 1;
    }
    utils::rayon::current_num_threads().next_power_of_two() * 2
}