// 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 math::FieldElement;
use utils::collections::Vec;

// AUXILIARY TRACE SEGMENT RANDOMNESS
// ================================================================================================

/// Random elements used in construction of auxiliary trace segments.
///
/// These elements are generated by the
/// [Air::get_aux_trace_segment_random_elements()](crate::Air::get_aux_trace_segment_random_elements)
/// function for each auxiliary trace segment. In the interactive version of the protocol, the
/// verifier draws these elements uniformly at random from the extension field of the protocol
/// after the prover commits to a previous trace segment.
#[derive(Debug, Clone)]
pub struct AuxTraceRandElements<E: FieldElement>(Vec<Vec<E>>);

impl<E: FieldElement> AuxTraceRandElements<E> {
    /// Instantiates and returns an empty set of random elements.
    pub fn new() -> Self {
        Self(Vec::new())
    }

    /// Returns a list of random elements for an auxiliary segment with the specified index.
    pub fn get_segment_elements(&self, aux_segment_idx: usize) -> &[E] {
        &self.0[aux_segment_idx]
    }

    /// Adds random elements for a new auxiliary segment to this set of random elements.
    pub fn add_segment_elements(&mut self, rand_elements: Vec<E>) {
        self.0.push(rand_elements);
    }
}

impl<E: FieldElement> Default for AuxTraceRandElements<E> {
    fn default() -> Self {
        Self::new()
    }
}

// CONSTRAINT COMPOSITION COEFFICIENTS
// ================================================================================================
/// Coefficients used in construction of constraint composition polynomial.
///
/// These coefficients are created by the
/// [Air::get_constraint_composition_coefficients()](crate::Air::get_constraint_composition_coefficients)
/// function. In the interactive version of the protocol, the verifier draws these coefficients
/// uniformly at random from the extension field of the protocol.
///
/// There is one coefficient for each constraint so that we can compute a random linear
/// combination of constraints as:
/// $$
/// \sum_{i = 0}^k{\alpha_i \cdot C_i(x)}
/// $$
/// where:
/// * $\alpha_i$ is the coefficient for the $i$th constraint.
/// * $C_i(x)$ is an evaluation of the $i$th constraint at $x$.
///
/// The coefficients are separated into two lists: one for transition constraints and another one
/// for boundary constraints. This separation is done for convenience only.
#[derive(Debug, Clone)]
pub struct ConstraintCompositionCoefficients<E: FieldElement> {
    pub transition: Vec<E>,
    pub boundary: Vec<E>,
}

// DEEP COMPOSITION COEFFICIENTS
// ================================================================================================
/// Coefficients used in construction of DEEP composition polynomial.
///
/// These coefficients are created by the
/// [Air::get_deep_composition_coefficients()](crate::Air::get_deep_composition_coefficients)
/// function. In the interactive version of the protocol, the verifier draws these coefficients
/// uniformly at random from the extension field of the protocol.
///
/// The coefficients are used in computing the DEEP composition polynomial as:
/// $$
/// Y(x) = \sum_{i=0}^k{(
///     \alpha_i \cdot (\frac{T_i(x) - T_i(z)}{x - z} +
///     \frac{T_i(x) - T_i(z \cdot g)}{x - z \cdot g})
/// )} + \sum_{j=0}^m{\beta_j \cdot \frac{H_j(x) - H_j(z)}{x - z}}
/// $$
/// where:
/// * $z$ is an out-of-domain point drawn randomly from the entire field. In the interactive
///   version of the protocol, $z$ is provided by the verifier.
/// * $g$ is the generator of the trace domain. This is the $n$th root of unity where
///   $n$ is the length of the execution trace.
/// * $T_i(x)$ is an evaluation of the $i$th trace polynomial at $x$, and $k$ is the total
///   number of trace polynomials (which is equal to the width of the execution trace).
/// * $H_i(x)$ is an evaluation of the $j$th constraint composition column polynomial at $x$,
///   and $m$ is the total number of column polynomials.
/// * $\alpha_i$ is a composition coefficient for the $i$th trace polynomial.
/// * $\beta_j$ is a composition coefficient for the $j$th constraint column polynomial.
///
/// The soundness of the resulting protocol with batching as above is given in Theorem 8 in
/// https://eprint.iacr.org/2022/1216 and it relies on two points:
///
/// 1. The evaluation proofs for each trace polynomial at $z$ and $g \cdot z$ can be batched using
/// the non-normalized Lagrange kernel over the set $\{z, g \cdot z\}$. This, however, requires
/// that the FRI protocol is run with rate $\rho^{+} := \frac{\kappa + 2}{\nu}$ where $\kappa$ and
/// $\nu$ are the length of the execution trace and the LDE domain size, respectively.
/// 2. The resulting $Y(x)$ do not need to be degree adjusted but the soundness error of the
/// protocol needs to be updated. For most combinations of batching parameters, this leads to a
/// negligible increase in soundness error. The formula for the updated error can be found in
/// Theorem 8 of https://eprint.iacr.org/2022/1216.
#[derive(Debug, Clone)]
pub struct DeepCompositionCoefficients<E: FieldElement> {
    /// Trace polynomial composition coefficients $\alpha_i$.
    pub trace: Vec<E>,
    /// Constraint column polynomial composition coefficients $\beta_j$.
    pub constraints: Vec<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.

use crate::{air::TransitionConstraintDegree, ProofOptions, TraceInfo};
use core::cmp;
use math::StarkField;
use utils::collections::Vec;

// AIR CONTEXT
// ================================================================================================
/// STARK parameters and trace properties for a specific execution of a computation.
#[derive(Clone, PartialEq, Eq)]
pub struct AirContext<B: StarkField> {
    pub(super) options: ProofOptions,
    pub(super) trace_info: TraceInfo,
    pub(super) main_transition_constraint_degrees: Vec<TransitionConstraintDegree>,
    pub(super) aux_transition_constraint_degrees: Vec<TransitionConstraintDegree>,
    pub(super) num_main_assertions: usize,
    pub(super) num_aux_assertions: usize,
    pub(super) ce_blowup_factor: usize,
    pub(super) trace_domain_generator: B,
    pub(super) lde_domain_generator: B,
    pub(super) num_transition_exemptions: usize,
}

impl<B: StarkField> AirContext<B> {
    // CONSTRUCTORS
    // --------------------------------------------------------------------------------------------
    /// Returns a new instance of [AirContext] instantiated for computations which require a single
    /// execution trace segment.
    ///
    /// The list of transition constraint degrees defines the total number of transition
    /// constraints and their expected degrees. Constraint evaluations computed by
    /// [Air::evaluate_transition()](crate::Air::evaluate_transition) function are expected to be
    /// in the order defined by this list.
    ///
    /// # Panics
    /// Panics if
    /// * `transition_constraint_degrees` is an empty vector.
    /// * `num_assertions` is zero.
    /// * Blowup factor specified by the provided `options` is too small to accommodate degrees
    ///   of the specified transition constraints.
    /// * `trace_info` describes a multi-segment execution trace.
    pub fn new(
        trace_info: TraceInfo,
        transition_constraint_degrees: Vec<TransitionConstraintDegree>,
        num_assertions: usize,
        options: ProofOptions,
    ) -> Self {
        assert!(
            !trace_info.is_multi_segment(),
            "provided trace info describes a multi-segment execution trace"
        );
        Self::new_multi_segment(
            trace_info,
            transition_constraint_degrees,
            Vec::new(),
            num_assertions,
            0,
            options,
        )
    }

    /// Returns a new instance of [AirContext] instantiated for computations which require multiple
    /// execution trace segments.
    ///
    /// The lists of transition constraint degrees defines the total number of transition
    /// constraints and their expected degrees. Constraint evaluations computed by
    /// [Air::evaluate_transition()](crate::Air::evaluate_transition) function are expected to be
    /// in the order defined by `main_transition_constraint_degrees` list. Constraint evaluations
    /// computed by [Air::evaluate_aux_transition()](crate::Air::evaluate_aux_transition) function
    /// are expected to be in the order defined by `aux_transition_constraint_degrees` list.
    ///
    /// # Panics
    /// Panics if
    /// * `main_transition_constraint_degrees` is an empty vector.
    /// * `num_main_assertions` is zero.
    /// * `trace_info.is_multi_segment() == true` but:
    ///   - `aux_transition_constraint_degrees` is an empty vector.
    ///   - `num_aux_assertions` is zero.
    /// * `trace_info.is_multi_segment() == false` but:
    ///   - `aux_transition_constraint_degrees` is a non-empty vector.
    ///   - `num_aux_assertions` is greater than zero.
    /// * Blowup factor specified by the provided `options` is too small to accommodate degrees
    ///   of the specified transition constraints.
    pub fn new_multi_segment(
        trace_info: TraceInfo,
        main_transition_constraint_degrees: Vec<TransitionConstraintDegree>,
        aux_transition_constraint_degrees: Vec<TransitionConstraintDegree>,
        num_main_assertions: usize,
        num_aux_assertions: usize,
        options: ProofOptions,
    ) -> Self {
        assert!(
            !main_transition_constraint_degrees.is_empty(),
            "at least one transition constraint degree must be specified"
        );
        assert!(
            num_main_assertions > 0,
            "at least one assertion must be specified"
        );

        if trace_info.is_multi_segment() {
            assert!(
                !aux_transition_constraint_degrees.is_empty(),
                "at least one transition constraint degree must be specified for auxiliary trace segments"
            );
            assert!(
                num_aux_assertions > 0,
                "at least one assertion must be specified against auxiliary trace segments"
            );
        } else {
            assert!(
                aux_transition_constraint_degrees.is_empty(),
                "auxiliary transition constraint degrees specified for a single-segment trace"
            );
            assert!(
                num_aux_assertions == 0,
                "auxiliary assertions specified for a single-segment trace"
            );
        }

        // determine minimum blowup factor needed to evaluate transition constraints by taking
        // the blowup factor of the highest degree constraint
        let mut ce_blowup_factor = 0;
        for degree in main_transition_constraint_degrees.iter() {
            if degree.min_blowup_factor() > ce_blowup_factor {
                ce_blowup_factor = degree.min_blowup_factor();
            }
        }

        for degree in aux_transition_constraint_degrees.iter() {
            if degree.min_blowup_factor() > ce_blowup_factor {
                ce_blowup_factor = degree.min_blowup_factor();
            }
        }

        assert!(
            options.blowup_factor() >= ce_blowup_factor,
            "blowup factor too small; expected at least {}, but was {}",
            ce_blowup_factor,
            options.blowup_factor()
        );

        let trace_length = trace_info.length();
        let lde_domain_size = trace_length * options.blowup_factor();

        AirContext {
            options,
            trace_info,
            main_transition_constraint_degrees,
            aux_transition_constraint_degrees,
            num_main_assertions,
            num_aux_assertions,
            ce_blowup_factor,
            trace_domain_generator: B::get_root_of_unity(trace_length.ilog2()),
            lde_domain_generator: B::get_root_of_unity(lde_domain_size.ilog2()),
            num_transition_exemptions: 1,
        }
    }

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

    /// Returns length of the execution trace for an instance of a computation.
    ///
    // This is guaranteed to be a power of two greater than or equal to 8.
    pub fn trace_len(&self) -> usize {
        self.trace_info.length()
    }

    /// Returns degree of trace polynomials for an instance of a computation.
    ///
    /// The degree is always `trace_length` - 1.
    pub fn trace_poly_degree(&self) -> usize {
        self.trace_info.length() - 1
    }

    /// Returns size of the constraint evaluation domain.
    ///
    /// This is guaranteed to be a power of two, and is equal to `trace_length * ce_blowup_factor`.
    pub fn ce_domain_size(&self) -> usize {
        self.trace_info.length() * self.ce_blowup_factor
    }

    /// Returns the size of the low-degree extension domain.
    ///
    /// This is guaranteed to be a power of two, and is equal to `trace_length * lde_blowup_factor`.
    pub fn lde_domain_size(&self) -> usize {
        self.trace_info.length() * self.options.blowup_factor()
    }

    /// Returns the number of transition constraints for a computation.
    ///
    /// The number of transition constraints is defined by the total number of transition
    /// constraint degree descriptors (for both the main and the auxiliary trace constraints).
    /// This number is used to determine how many transition constraint coefficients need to be
    /// generated for merging transition constraints into a composition polynomial.
    pub fn num_transition_constraints(&self) -> usize {
        self.main_transition_constraint_degrees.len() + self.aux_transition_constraint_degrees.len()
    }

    /// Returns the number of transition constraints placed against the main trace segment.
    pub fn num_main_transition_constraints(&self) -> usize {
        self.main_transition_constraint_degrees.len()
    }

    /// Returns the number of transition constraints placed against all auxiliary trace segments.
    pub fn num_aux_transition_constraints(&self) -> usize {
        self.aux_transition_constraint_degrees.len()
    }

    /// Returns the total number of assertions defined for a computation.
    ///
    /// The number of assertions consists of the assertions placed against the main segment of an
    /// execution trace as well as assertions placed against all auxiliary trace segments.
    pub fn num_assertions(&self) -> usize {
        self.num_main_assertions + self.num_aux_assertions
    }

    /// Returns the number of rows at the end of an execution trace to which transition constraints
    /// do not apply.
    ///
    /// This is guaranteed to be at least 1 (which is the default value), but could be greater.
    /// The maximum number of exemptions is determined by a combination of transition constraint
    /// degrees and blowup factor specified for the computation.
    pub fn num_transition_exemptions(&self) -> usize {
        self.num_transition_exemptions
    }

    /// Returns the number of columns needed to store the constraint composition polynomial.
    ///
    /// This is the maximum of:
    /// 1. The maximum evaluation degree over all transition constraints minus the degree
    /// of the transition constraint divisor divided by trace length.
    /// 2. `1`, because the constraint composition polynomial requires at least one column.
    ///
    /// Since the degree of a constraint `C(x)` can be well approximated by
    /// `[constraint.base + constraint.cycles.len()] * [trace_length - 1]` the degree of the
    /// constraint composition polynomial can be approximated by:
    /// `([constraint.base + constraint.cycles.len()] * [trace_length - 1] - [trace_length - n])`
    /// where `constraint` is the constraint attaining the maximum and `n` is the number of
    /// exemption points.
    /// In the case `n = 1`, the expression simplifies to:
    /// `[constraint.base + constraint.cycles.len() - 1] * [trace_length - 1]`
    /// Thus, if each column is of length `trace_length`, we would need
    /// `[constraint.base + constraint.cycles.len() - 1]` columns to store the coefficients of
    /// the constraint composition polynomial.
    /// This means that if the highest constraint degree is equal to `5`, the constraint
    /// composition polynomial will require four columns and if the highest constraint degree is
    /// equal to `7`, it will require six columns to store.
    pub fn num_constraint_composition_columns(&self) -> usize {
        let mut highest_constraint_degree = 0_usize;
        for degree in self
            .main_transition_constraint_degrees
            .iter()
            .chain(self.aux_transition_constraint_degrees.iter())
        {
            let eval_degree = degree.get_evaluation_degree(self.trace_len());
            if eval_degree > highest_constraint_degree {
                highest_constraint_degree = eval_degree
            }
        }
        let trace_length = self.trace_len();
        let transition_divisior_degree = trace_length - self.num_transition_exemptions();

        // we use the identity: ceil(a/b) = (a + b - 1)/b
        let num_constraint_col =
            (highest_constraint_degree - transition_divisior_degree + trace_length - 1)
                / trace_length;

        cmp::max(num_constraint_col, 1)
    }

    // DATA MUTATORS
    // --------------------------------------------------------------------------------------------

    /// Sets the number of transition exemptions for this context.
    ///
    /// # Panics
    /// Panics if:
    /// * The number of exemptions is zero.
    /// * The number of exemptions exceeds half of the trace length.
    /// * Given the combination of transition constraints degrees and the blowup factor in this
    ///   context, the number of exemptions is too larger for a valid computation of the constraint
    ///   composition polynomial.
    pub fn set_num_transition_exemptions(mut self, n: usize) -> Self {
        assert!(
            n > 0,
            "number of transition exemptions must be greater than zero"
        );
        // exemptions which are for more than half the trace plus one are probably a mistake
        assert!(
            n <= self.trace_len() / 2 + 1,
            "number of transition exemptions cannot exceed {}, but was {}",
            self.trace_len() / 2 + 1,
            n
        );
        // make sure the composition polynomial can be computed correctly with the specified
        // number of exemptions.
        // The `ce_blowup` factor puts a ceiling on the maximal degree of a constraint composition
        // polynomial we can accomodate. On the other hand, adding exemption points reduces the
        // degree of the divisor which results in an increase of the resulting constraint composition
        // polynomial.Thus we need to check that the number of exemption points is not too large
        // given the above.
        for degree in self
            .main_transition_constraint_degrees
            .iter()
            .chain(self.aux_transition_constraint_degrees.iter())
        {
            let eval_degree = degree.get_evaluation_degree(self.trace_len());
            let max_constraint_composition_degree = self.ce_domain_size() - 1;
            let max_exemptions = max_constraint_composition_degree + self.trace_len() - eval_degree;
            assert!(
                n <= max_exemptions,
                "number of transition exemptions cannot exceed: {max_exemptions}, but was {n}"
            )
        }

        self.num_transition_exemptions = n;
        self
    }
}


// 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::air::Assertion;
use core::fmt::{Display, Formatter};
use math::{FieldElement, StarkField};
use utils::collections::Vec;

// CONSTRAINT DIVISOR
// ================================================================================================
/// The denominator portion of boundary and transition constraints.
///
/// A divisor is described by a combination of a sparse polynomial, which describes the numerator
/// of the divisor and a set of exemption points, which describe the denominator of the divisor.
/// The numerator polynomial is described as multiplication of tuples where each tuple encodes
/// an expression $(x^a - b)$. The exemption points encode expressions $(x - a)$.
///
/// For example divisor $(x^a - 1) \cdot (x^b - 2) / (x - 3)$ can be represented as:
/// numerator: `[(a, 1), (b, 2)]`, exemptions: `[3]`.
///
/// A divisor cannot be instantiated directly, and instead must be created either for an
/// [Assertion] or for a transition constraint.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ConstraintDivisor<B: StarkField> {
    pub(super) numerator: Vec<(usize, B)>,
    pub(super) exemptions: Vec<B>,
}

impl<B: StarkField> ConstraintDivisor<B> {
    // CONSTRUCTORS
    // --------------------------------------------------------------------------------------------

    /// Returns a new divisor instantiated from the provided parameters.
    fn new(numerator: Vec<(usize, B)>, exemptions: Vec<B>) -> Self {
        ConstraintDivisor {
            numerator,
            exemptions,
        }
    }

    /// Builds a divisor for transition constraints.
    ///
    /// For transition constraints, the divisor polynomial $z(x)$ is always the same:
    ///
    /// $$
    /// z(x) = \frac{x^n - 1}{ \prod_{i=1}^k (x - g^{n-i})}
    /// $$
    ///
    /// where, $n$ is the length of the execution trace, $g$ is the generator of the trace
    /// domain, and $k$ is the number of exemption points. The default value for $k$ is $1$.
    ///
    /// The above divisor specifies that transition constraints must hold on all steps of the
    /// execution trace except for the last $k$ steps.
    pub fn from_transition(trace_length: usize, num_exemptions: usize) -> Self {
        assert!(
            num_exemptions > 0,
            "invalid number of transition exemptions: must be greater than zero"
        );
        let exemptions = (trace_length - num_exemptions..trace_length)
            .map(|step| get_trace_domain_value_at::<B>(trace_length, step))
            .collect();
        Self::new(vec![(trace_length, B::ONE)], exemptions)
    }

    /// Builds a divisor for a boundary constraint described by the assertion.
    ///
    /// For boundary constraints, the divisor polynomial is defined as:
    ///
    /// $$
    /// z(x) = x^k - g^{a \cdot k}
    /// $$
    ///
    /// where $g$ is the generator of the trace domain, $k$ is the number of asserted steps, and
    /// $a$ is the step offset in the trace domain. Specifically:
    /// * For an assertion against a single step, the polynomial is $(x - g^a)$, where $a$ is the
    ///   step on which the assertion should hold.
    /// * For an assertion against a sequence of steps which fall on powers of two, it is
    ///   $(x^k - 1)$ where $k$ is the number of asserted steps.
    /// * For assertions against a sequence of steps which repeat with a period that is a power
    ///   of two but don't fall exactly on steps which are powers of two (e.g. 1, 9, 17, ... )
    ///   it is $(x^k - g^{a \cdot k})$, where $a$ is the number of steps by which the assertion steps
    ///   deviate from a power of two, and $k$ is the number of asserted steps. This is equivalent to
    ///   $(x - g^a) \cdot (x - g^{a + j}) \cdot (x - g^{a + 2 \cdot j}) ... (x - g^{a + (k  - 1) \cdot j})$,
    ///   where $j$ is the length of interval between asserted steps (e.g. 8).
    ///
    /// # Panics
    /// Panics of the specified `trace_length` is inconsistent with the specified `assertion`.
    pub fn from_assertion<E>(assertion: &Assertion<E>, trace_length: usize) -> Self
    where
        E: FieldElement<BaseField = B>,
    {
        let num_steps = assertion.get_num_steps(trace_length);
        if assertion.first_step == 0 {
            Self::new(vec![(num_steps, B::ONE)], vec![])
        } else {
            let trace_offset = num_steps * assertion.first_step;
            let offset = get_trace_domain_value_at::<B>(trace_length, trace_offset);
            Self::new(vec![(num_steps, offset)], vec![])
        }
    }

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

    /// Returns the numerator portion of this constraint divisor.
    pub fn numerator(&self) -> &[(usize, B)] {
        &self.numerator
    }

    /// Returns exemption points (the denominator portion) of this constraints divisor.
    pub fn exemptions(&self) -> &[B] {
        &self.exemptions
    }

    /// Returns the degree of the divisor polynomial
    pub fn degree(&self) -> usize {
        let numerator_degree = self
            .numerator
            .iter()
            .fold(0, |degree, term| degree + term.0);
        let denominator_degree = self.exemptions.len();
        numerator_degree - denominator_degree
    }

    // EVALUATORS
    // --------------------------------------------------------------------------------------------
    /// Evaluates the divisor polynomial at the provided `x` coordinate.
    pub fn evaluate_at<E: FieldElement<BaseField = B>>(&self, x: E) -> E {
        // compute the numerator value
        let mut numerator = E::ONE;
        for (degree, constant) in self.numerator.iter() {
            let v = x.exp((*degree as u32).into());
            let v = v - E::from(*constant);
            numerator *= v;
        }

        // compute the denominator value
        let denominator = self.evaluate_exemptions_at(x);

        numerator / denominator
    }

    /// Evaluates the denominator of this divisor (the exemption points) at the provided `x`
    /// coordinate.
    #[inline(always)]
    pub fn evaluate_exemptions_at<E: FieldElement<BaseField = B>>(&self, x: E) -> E {
        self.exemptions
            .iter()
            .fold(E::ONE, |r, &e| r * (x - E::from(e)))
    }
}

impl<B: StarkField> Display for ConstraintDivisor<B> {
    fn fmt(&self, f: &mut Formatter) -> core::fmt::Result {
        for (degree, offset) in self.numerator.iter() {
            write!(f, "(x^{degree} - {offset})")?;
        }
        if !self.exemptions.is_empty() {
            write!(f, " / ")?;
            for x in self.exemptions.iter() {
                write!(f, "(x - {x})")?;
            }
        }
        Ok(())
    }
}

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

/// Returns g^step, where g is the generator of trace domain.
fn get_trace_domain_value_at<B: StarkField>(trace_length: usize, step: usize) -> B {
    debug_assert!(
        step < trace_length,
        "step must be in the trace domain [0, {trace_length})"
    );
    let g = B::get_root_of_unity(trace_length.ilog2());
    g.exp((step as u64).into())
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use math::{fields::f128::BaseElement, polynom};

    #[test]
    fn constraint_divisor_degree() {
        // single term numerator
        let div = ConstraintDivisor::new(vec![(4, BaseElement::ONE)], vec![]);
        assert_eq!(4, div.degree());

        // multi-term numerator
        let div = ConstraintDivisor::new(
            vec![
                (4, BaseElement::ONE),
                (2, BaseElement::new(2)),
                (3, BaseElement::new(3)),
            ],
            vec![],
        );
        assert_eq!(9, div.degree());

        // multi-term numerator with exemption points
        let div = ConstraintDivisor::new(
            vec![
                (4, BaseElement::ONE),
                (2, BaseElement::new(2)),
                (3, BaseElement::new(3)),
            ],
            vec![BaseElement::ONE, BaseElement::new(2)],
        );
        assert_eq!(7, div.degree());
    }

    #[test]
    fn constraint_divisor_evaluation() {
        // single term numerator: (x^4 - 1)
        let div = ConstraintDivisor::new(vec![(4, BaseElement::ONE)], vec![]);
        assert_eq!(BaseElement::new(15), div.evaluate_at(BaseElement::new(2)));

        // multi-term numerator: (x^4 - 1) * (x^2 - 2) * (x^3 - 3)
        let div = ConstraintDivisor::new(
            vec![
                (4, BaseElement::ONE),
                (2, BaseElement::new(2)),
                (3, BaseElement::new(3)),
            ],
            vec![],
        );
        let expected = BaseElement::new(15) * BaseElement::new(2) * BaseElement::new(5);
        assert_eq!(expected, div.evaluate_at(BaseElement::new(2)));

        // multi-term numerator with exemption points:
        // (x^4 - 1) * (x^2 - 2) * (x^3 - 3) / ((x - 1) * (x - 2))
        let div = ConstraintDivisor::new(
            vec![
                (4, BaseElement::ONE),
                (2, BaseElement::new(2)),
                (3, BaseElement::new(3)),
            ],
            vec![BaseElement::ONE, BaseElement::new(2)],
        );
        let expected = BaseElement::new(255) * BaseElement::new(14) * BaseElement::new(61)
            / BaseElement::new(6);
        assert_eq!(expected, div.evaluate_at(BaseElement::new(4)));
    }

    #[test]
    fn constraint_divisor_equivalence() {
        let n = 8_usize;
        let g = BaseElement::get_root_of_unity(n.trailing_zeros());
        let k = 4_u32;
        let j = n as u32 / k;

        // ----- periodic assertion divisor, no offset --------------------------------------------

        // create a divisor for assertion which repeats every 2 steps starting at step 0
        let assertion = Assertion::periodic(0, 0, j as usize, BaseElement::ONE);
        let divisor = ConstraintDivisor::from_assertion(&assertion, n);

        // z(x) = x^4 - 1 = (x - 1) * (x - g^2) * (x - g^4) * (x - g^6)
        let poly = polynom::mul(
            &polynom::mul(
                &[-BaseElement::ONE, BaseElement::ONE],
                &[-g.exp(j.into()), BaseElement::ONE],
            ),
            &polynom::mul(
                &[-g.exp((2 * j).into()), BaseElement::ONE],
                &[-g.exp((3 * j).into()), BaseElement::ONE],
            ),
        );

        for i in 0..n {
            let expected = polynom::eval(&poly, g.exp((i as u32).into()));
            let actual = divisor.evaluate_at(g.exp((i as u32).into()));
            assert_eq!(expected, actual);
            if i % (j as usize) == 0 {
                assert_eq!(BaseElement::ZERO, actual);
            }
        }

        // ----- periodic assertion divisor, with offset ------------------------------------------

        // create a divisor for assertion which repeats every 2 steps starting at step 1
        let offset = 1_u32;
        let assertion = Assertion::periodic(0, offset as usize, j as usize, BaseElement::ONE);
        let divisor = ConstraintDivisor::from_assertion(&assertion, n);
        assert_eq!(
            ConstraintDivisor::new(vec![(k as usize, g.exp(k.into()))], vec![]),
            divisor
        );

        // z(x) = x^4 - g^4 = (x - g) * (x - g^3) * (x - g^5) * (x - g^7)
        let poly = polynom::mul(
            &polynom::mul(
                &[-g.exp(offset.into()), BaseElement::ONE],
                &[-g.exp((offset + j).into()), BaseElement::ONE],
            ),
            &polynom::mul(
                &[-g.exp((offset + 2 * j).into()), BaseElement::ONE],
                &[-g.exp((offset + 3 * j).into()), BaseElement::ONE],
            ),
        );

        for i in 0..n {
            let expected = polynom::eval(&poly, g.exp((i as u32).into()));
            let actual = divisor.evaluate_at(g.exp((i as u32).into()));
            assert_eq!(expected, actual);
            if i % (j as usize) == offset as usize {
                assert_eq!(BaseElement::ZERO, actual);
            }
        }

        // create a divisor for assertion which repeats every 4 steps starting at step 3
        let offset = 3_u32;
        let k = 2_u32;
        let j = n as u32 / k;
        let assertion = Assertion::periodic(0, offset as usize, j as usize, BaseElement::ONE);
        let divisor = ConstraintDivisor::from_assertion(&assertion, n);
        assert_eq!(
            ConstraintDivisor::new(vec![(k as usize, g.exp((offset * k).into()))], vec![]),
            divisor
        );

        // z(x) = x^2 - g^6 = (x - g^3) * (x - g^7)
        let poly = polynom::mul(
            &[-g.exp(offset.into()), BaseElement::ONE],
            &[-g.exp((offset + j).into()), BaseElement::ONE],
        );

        for i in 0..n {
            let expected = polynom::eval(&poly, g.exp((i as u32).into()));
            let actual = divisor.evaluate_at(g.exp((i as u32).into()));
            assert_eq!(expected, actual);
            if i % (j as usize) == offset as usize {
                assert_eq!(BaseElement::ZERO, 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 crate::ProofOptions;
use crypto::{RandomCoin, RandomCoinError};
use math::{fft, ExtensibleField, ExtensionOf, FieldElement, StarkField, ToElements};
use utils::collections::{BTreeMap, Vec};

mod trace_info;
pub use trace_info::{TraceInfo, TraceLayout};

mod context;
pub use context::AirContext;

mod assertions;
pub use assertions::Assertion;

mod boundary;
pub use boundary::{BoundaryConstraint, BoundaryConstraintGroup, BoundaryConstraints};

mod transition;
pub use transition::{EvaluationFrame, TransitionConstraintDegree, TransitionConstraints};

mod coefficients;
pub use coefficients::{
    AuxTraceRandElements, ConstraintCompositionCoefficients, DeepCompositionCoefficients,
};

mod divisor;
pub use divisor::ConstraintDivisor;

#[cfg(test)]
mod tests;

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

const MIN_CYCLE_LENGTH: usize = 2;

// AIR TRAIT
// ================================================================================================
/// Describes algebraic intermediate representation of a computation.
///
/// To describe AIR for a given computation, you'll need to implement the `Air` trait which
/// involves the following:
///
/// 1. Define base field for your computation via the [Air::BaseField] associated type (see
///    [math::fields] for available field options).
/// 2. Define a set of public inputs which are required for your computation via the
///    [Air::PublicInputs] associated type.
/// 3. Implement [Air::new()] function. As a part of this function you should create a
///    [AirContext] struct which takes degrees for all transition constraints as one of
///    the constructor parameters.
/// 4. Implement [Air::context()] method which should return a reference to the
///    [AirContext] struct created in [Air::new()] function.
/// 5. Implement [Air::evaluate_transition()] method which should evaluate
///    [transition constraints](#transition-constraints) over a given evaluation frame.
/// 6. Implement [Air::get_assertions()] method which should return a vector of
///    [assertions](#trace-assertions) for a given instance of your computation.
/// 7. If your computation requires [periodic values](#periodic-values), you can also override
///    the default [Air::get_periodic_column_values()] method.
///
/// If your computation uses [Randomized AIR](#randomized-air), you will also need to override
/// [Air::evaluate_aux_transition()] and [Air::get_aux_assertions()] methods.
///
/// ### Transition constraints
/// Transition constraints define algebraic relations between two consecutive steps of a
/// computation. In Winterfell, transition constraints are evaluated inside
/// [Air::evaluate_transition()] function which takes the following parameters:
///
/// - [EvaluationFrame] which contains vectors with current and next states of the
///   computation.
/// - A list of periodic values. When periodic columns are defined for a computation,
///   this will contain values of periodic columns at the current step of the computation.
///   Otherwise, this will be an empty list.
/// - A mutable `result` slice. This is the slice where constraint evaluations should be
///   written to. The length of this slice will be equal to the number of transition
///   constraints defined for the computation.
///
/// The constraints are considered to be satisfied if and only if, after the function returns,
/// the `result` slice contains all zeros. In general, it is important for the transition
/// constraint evaluation function to work as follows:
///
/// * For all valid transitions between consecutive computation steps, transition constraints
///   should evaluation to all zeros.
/// * For any invalid transition, at least one constraint must evaluate to a non-zero value.
///
/// **Note:** since transition constraints define algebraic relations, they should be
/// described using only algebraic operations: additions, subtractions, and multiplications
/// (divisions can be emulated using inverse of multiplication).
///
/// ### Constraint degrees
/// One of the main factors impacting proof generation time and proof size is the maximum degree
/// of transition constraints. The higher is this degree, the larger our blowup factor needs to be.
/// Usually, we want to keep this degree as low as possible - e.g. under 4 or 8. To accurately
/// describe degrees of your transition constraints, keep the following in mind:
///
/// * All trace columns have degree `1`.
/// * When multiplying trace columns together, the degree increases by `1`. For example, if our
///   constraint involves multiplication of two columns, the degree of this constraint will be
///   `2`. We can describe this constraint using [TransitionConstraintDegree] struct as follows:
///   `TransitionConstraintDegree::new(2)`.
/// * Degrees of periodic columns depend on the length of their cycles, but in most cases, these
///   degrees are very close to `1`.
/// * To describe a degree of a constraint involving multiplication of trace columns and
///   periodic columns, use the [TransitionConstraintDegree::with_cycles()] constructor. For
///   example, if our constraint involves multiplication of one trace column and one periodic
///   column with a cycle of 32 steps, the degree can be described as:
///   `TransitionConstraintDegree::with_cycles(1, vec![32])`.
///
/// In general, multiplications should be used judiciously - though, there are ways to ease this
/// restriction a bit at the expense of wider execution trace.
///
/// ### Trace assertions
/// Assertions are used to specify that a valid execution trace of a computation must contain
/// certain values in certain cells. They are frequently used to tie public inputs to a specific
/// execution trace, but can be used to constrain a computation in other ways as well.
/// Internally within Winterfell, assertions are converted into *boundary constraints*.
///
/// To define assertions for your computation, you'll need to implement [Air::get_assertions()]
/// function which should return a vector of [Assertion] structs. Every computation must have at
/// least one assertion. Assertions can be of the following types:
///
/// * A single assertion - such assertion specifies that a single cell of an execution trace must
///   be equal to a specific value. For example: *value in column 0, at step 0, must be equal
///   to 1*.
/// * A periodic assertion - such assertion specifies that values in a given column at specified
///   intervals should be equal to some value. For example: *values in column 0, at steps 0, 8,
///   16, 24 etc. must be equal to 2*.
/// * A sequence assertion - such assertion specifies that values in a given column at specific
///   intervals must be equal to a sequence of provided values. For example: *values in column 0,
///   at step 0 must be equal to 1, at step 8 must be equal to 2, at step 16 must be equal to 3
///   etc.*
///
/// ### Periodic values
/// Sometimes, it may be useful to define a column in an execution trace which contains a set of
/// repeating values. For example, let's say we have a column which contains value 1 on every
/// 4th step, and 0 otherwise. Such a column can be described with a simple periodic sequence of
/// `[1, 0, 0, 0]`.
///
/// To define such columns for your computation, you can override
/// [Air::get_periodic_column_values()] method. The values of the periodic columns at a given
/// step of the computation will be supplied to the [Air::evaluate_transition()] method via the
/// `periodic_values` parameter.
///
/// ### Randomized AIR
/// Randomized AIR is a powerful extension of AIR which enables, among other things, multiset and
/// permutation checks similar to the ones available in PLONKish systems. These, in turn, allow
/// efficient descriptions of "non-local" constraints which can be used to build such components
/// as efficient range checks, random access memory, and many others.
///
/// With Randomized AIR, construction of the execution trace is split into multiple stages. During
/// the first stage, the *main trace segment* is built in a manner similar to how the trace is
/// built for regular AIR. In the subsequent stages, *auxiliary trace segments* are built. When
/// building auxiliary trace segments, the prover has access to extra randomness sent by the
/// verifier (in the non-interactive version of the protocol, this randomness is derived from the
/// previous trace segment commitments). Currently, the number of auxiliary trace segments is
/// limited to one.
///
/// To describe Randomized AIR, you will need to do the following when implementing the [Air]
/// trait:
/// * The [AirContext] struct returned from [Air::context()] method must be instantiated using
///   [AirContext::new_multi_segment()] constructor. When building AIR context in this way, you
///   will need to provide a [TraceLayout] which describes the shape of a multi-segment execution
///   trace.
/// * Override [Air::evaluate_aux_transition()] method. This method is similar to the
///   [Air::evaluate_transition()] method but it also accepts two extra parameters:
///   `aux_evaluation_frame` and `aux_rand_elements`. These parameters are needed for evaluating
///   transition constraints over the auxiliary trace segments.
/// * Override [Air::get_aux_assertions()] method. This method is similar to the
///   [Air::get_assertions()] method, but it should return assertions against columns of the
///   auxiliary trace segments.
pub trait Air: Send + Sync {
    /// Base field for the computation described by this AIR. STARK protocol for this computation
    /// may be executed in the base field, or in an extension of the base fields as specified
    /// by [ProofOptions] struct.
    type BaseField: StarkField + ExtensibleField<2> + ExtensibleField<3>;

    /// A type defining shape of public inputs for the computation described by this protocol.
    /// This could be any type as long as it can be serialized into a sequence of field elements.
    type PublicInputs: ToElements<Self::BaseField>;

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

    /// Returns new instance of AIR for this computation instantiated from the provided parameters,
    /// which have the following meaning:
    /// - `trace_info` contains information about a concrete execution trace of the computation
    ///   described by this AIR, including trace width, trace length length, and optionally,
    ///   additional custom parameters in `meta` field.
    /// - `public_inputs` specifies public inputs for this instance of the computation.
    /// - `options` defines proof generation options such as blowup factor, hash function etc.
    ///   these options define security level of the proof and influence proof generation time.
    fn new(trace_info: TraceInfo, pub_inputs: Self::PublicInputs, options: ProofOptions) -> Self;

    /// Returns context for this instance of the computation.
    fn context(&self) -> &AirContext<Self::BaseField>;

    /// Evaluates transition constraints over the specified evaluation frame.
    ///
    /// The evaluations should be written into the `results` slice in the same order as the
    /// the order of transition constraint degree descriptors used to instantiate [AirContext]
    /// for this AIR. Thus, the length of the `result` slice will equal to the number of
    /// transition constraints defined for this computation.
    ///
    /// We define type `E` separately from `Self::BaseField` to allow evaluation of constraints
    /// over the out-of-domain evaluation frame, which may be defined over an extension field
    /// (when extension fields are used).
    fn evaluate_transition<E: FieldElement<BaseField = Self::BaseField>>(
        &self,
        frame: &EvaluationFrame<E>,
        periodic_values: &[E],
        result: &mut [E],
    );

    /// Returns a set of assertions against a concrete execution trace of this computation.
    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>>;

    // AUXILIARY TRACE CONSTRAINTS
    // --------------------------------------------------------------------------------------------

    /// Evaluates transition constraints over the specified evaluation frames for the main and
    /// auxiliary trace segments.
    ///
    /// The evaluations should be written into the `results` slice in the same order as the order
    /// of auxiliary transition constraint degree descriptors used to instantiate [AirContext] for
    /// this AIR. Thus, the length of the `result` slice will equal to the number of auxiliary
    /// transition constraints defined for this computation.
    ///
    /// The default implementation of this function panics. It must be overridden for AIRs
    /// describing computations which require multiple trace segments.
    ///
    /// The types for main and auxiliary trace evaluation frames are defined as follows:
    /// * When the entire protocol is executed in a prime field, types `F` and `E` are the same,
    ///   and thus, both the main and the auxiliary trace frames are defined over the base field.
    /// * When the protocol is executed in an extension field, the main trace frame is defined
    ///   over the base field, while the auxiliary trace frame is defined over the extension field.
    ///
    /// We define type `F` separately from `Self::BaseField` to allow evaluation of constraints
    /// over the out-of-domain evaluation frame, which may be defined over an extension field
    /// (when extension fields are used). The type bounds specified for this function allow the
    /// following:
    /// * `F` and `E` could be the same [StarkField] or extensions of the same [StarkField].
    /// * `F` and `E` could be the same field, because a field is always an extension of itself.
    /// * If `F` and `E` are different, then `E` must be an extension of `F`.
    #[allow(unused_variables)]
    fn evaluate_aux_transition<F, E>(
        &self,
        main_frame: &EvaluationFrame<F>,
        aux_frame: &EvaluationFrame<E>,
        periodic_values: &[F],
        aux_rand_elements: &AuxTraceRandElements<E>,
        result: &mut [E],
    ) where
        F: FieldElement<BaseField = Self::BaseField>,
        E: FieldElement<BaseField = Self::BaseField> + ExtensionOf<F>,
    {
        unimplemented!("evaluation of auxiliary transition constraints has not been implemented");
    }

    /// Returns a set of assertions placed against auxiliary trace segments.
    ///
    /// The default implementation of this function returns an empty vector. It should be
    /// overridden only if the computation relies on auxiliary trace segments. In such a case,
    /// the vector returned from this function must contain at least one assertion.
    ///
    /// The column index for assertions is expected to be zero-based across all auxiliary trace
    /// segments. That is, assertion against column 0, is an assertion against the first column
    /// of the auxiliary trace segments.
    ///
    /// When the protocol is executed using an extension field, auxiliary assertions are defined
    /// over the extension field. This is in contrast with the assertions returned from
    /// [get_assertions()](Air::get_assertions) function, which always returns assertions defined
    /// over the base field of the protocol.
    #[allow(unused_variables)]
    fn get_aux_assertions<E: FieldElement<BaseField = Self::BaseField>>(
        &self,
        aux_rand_elements: &AuxTraceRandElements<E>,
    ) -> Vec<Assertion<E>> {
        Vec::new()
    }

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

    /// Returns values for all periodic columns used in the computation.
    ///
    /// These values will be used to compute column values at specific states of the computation
    /// and passed in to the [evaluate_transition()](Air::evaluate_transition) method as
    /// `periodic_values` parameter.
    ///
    /// The default implementation of this method returns an empty vector. For computations which
    /// rely on periodic columns, this method should be overridden in the specialized
    /// implementation. Number of values for each periodic column must be a power of two.
    fn get_periodic_column_values(&self) -> Vec<Vec<Self::BaseField>> {
        Vec::new()
    }

    /// Returns polynomial for all periodic columns.
    ///
    /// These polynomials are interpolated from the values returned from the
    /// [get_periodic_column_values()](Air::get_periodic_column_values) method.
    fn get_periodic_column_polys(&self) -> Vec<Vec<Self::BaseField>> {
        // cache inverse twiddles for each cycle length so that we don't have to re-build them
        // for columns with identical cycle lengths
        let mut twiddle_map = BTreeMap::new();
        // iterate over all periodic columns and convert column values into polynomials
        self.get_periodic_column_values()
            .into_iter()
            .map(|mut column| {
                let cycle_length = column.len();
                assert!(
                    cycle_length >= MIN_CYCLE_LENGTH,
                    "number of values in a periodic column must be at least {MIN_CYCLE_LENGTH}, but was {cycle_length}"
                );
                assert!(
                    cycle_length.is_power_of_two(),
                    "number of values in a periodic column must be a power of two, but was {cycle_length}"
                );
                assert!(cycle_length <= self.trace_length(),
                    "number of values in a periodic column cannot exceed trace length {}, but was {}",
                    self.trace_length(),
                    cycle_length
                );

                // get twiddles for interpolation and interpolate values into a polynomial
                let inv_twiddles = twiddle_map
                    .entry(cycle_length)
                    .or_insert_with(|| fft::get_inv_twiddles::<Self::BaseField>(cycle_length));
                fft::interpolate_poly(&mut column, inv_twiddles);
                column
            })
            .collect()
    }

    /// Groups transition constraints together by their degree.
    ///
    /// This function also assigns composition coefficients to each constraint. These coefficients
    /// will be used to compute a random linear combination of transition constraints evaluations
    /// during constraint merging performed by [TransitionConstraintGroup::merge_evaluations()]
    /// function.
    fn get_transition_constraints<E: FieldElement<BaseField = Self::BaseField>>(
        &self,
        composition_coefficients: &[E],
    ) -> TransitionConstraints<E> {
        TransitionConstraints::new(self.context(), composition_coefficients)
    }

    /// Convert assertions returned from [get_assertions()](Air::get_assertions) and
    /// [get_aux_assertions()](Air::get_aux_assertions) methods into boundary constraints.
    ///
    /// This function also assigns composition coefficients to each constraint, and groups the
    /// constraints by their divisors. The coefficients will be used to compute random linear
    /// combination of boundary constraints during constraint merging.
    fn get_boundary_constraints<E: FieldElement<BaseField = Self::BaseField>>(
        &self,
        aux_rand_elements: &AuxTraceRandElements<E>,
        composition_coefficients: &[E],
    ) -> BoundaryConstraints<E> {
        BoundaryConstraints::new(
            self.context(),
            self.get_assertions(),
            self.get_aux_assertions(aux_rand_elements),
            composition_coefficients,
        )
    }

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

    /// Returns options which specify STARK protocol parameters for an instance of the computation
    /// described by this AIR.
    fn options(&self) -> &ProofOptions {
        &self.context().options
    }

    /// Returns info of the execution trace for an instance of the computation described by
    /// this AIR.
    fn trace_info(&self) -> &TraceInfo {
        &self.context().trace_info
    }

    /// Returns length of the execution trace for an instance of the computation described by
    /// this AIR.
    ///
    // This is guaranteed to be a power of two greater than or equal to 8.
    fn trace_length(&self) -> usize {
        self.context().trace_info.length()
    }

    /// Returns a description of how execution trace columns are arranged into segments for
    /// an instance of a computation described by this AIR.
    fn trace_layout(&self) -> &TraceLayout {
        self.context().trace_info.layout()
    }

    /// Returns degree of trace polynomials for an instance of the computation described by
    /// this AIR.
    ///
    /// The degree is always `trace_length` - 1.
    fn trace_poly_degree(&self) -> usize {
        self.context().trace_poly_degree()
    }

    /// Returns the generator of the trace domain for an instance of the computation described
    /// by this AIR.
    ///
    /// The generator is the $n$th root of unity where $n$ is the length of the execution trace.
    fn trace_domain_generator(&self) -> Self::BaseField {
        self.context().trace_domain_generator
    }

    /// Returns constraint evaluation domain blowup factor for the computation described by this
    /// AIR.
    ///
    /// The blowup factor is defined as the smallest power of two greater than or equal to the
    /// hightest transition constraint degree. For example, if the hightest transition
    /// constraint degree = 3, `ce_blowup_factor` will be set to 4.
    ///
    /// `ce_blowup_factor` is guaranteed to be smaller than or equal to the `lde_blowup_factor`.
    fn ce_blowup_factor(&self) -> usize {
        self.context().ce_blowup_factor
    }

    /// Returns size of the constraint evaluation domain.
    ///
    /// This is guaranteed to be a power of two, and is equal to `trace_length * ce_blowup_factor`.
    fn ce_domain_size(&self) -> usize {
        self.context().ce_domain_size()
    }

    /// Returns low-degree extension domain blowup factor for the computation described by this
    /// AIR. This is guaranteed to be a power of two, and is always either equal to or greater
    /// than ce_blowup_factor.
    fn lde_blowup_factor(&self) -> usize {
        self.context().options.blowup_factor()
    }

    /// Returns the size of the low-degree extension domain.
    ///
    /// This is guaranteed to be a power of two, and is equal to `trace_length * lde_blowup_factor`.
    fn lde_domain_size(&self) -> usize {
        self.context().lde_domain_size()
    }

    /// Returns the generator of the low-degree extension domain for an instance of the
    /// computation described by this AIR.
    ///
    /// The generator is the $n$th root of unity where $n$ is the size of the low-degree extension
    /// domain.
    fn lde_domain_generator(&self) -> Self::BaseField {
        self.context().lde_domain_generator
    }

    /// Returns the offset by which the domain for low-degree extension is shifted in relation
    /// to the execution trace domain.
    fn domain_offset(&self) -> Self::BaseField {
        self.context().options.domain_offset()
    }

    // TRACE SEGMENT RANDOMNESS
    // --------------------------------------------------------------------------------------------

    /// Returns a vector of field elements required for construction of an auxiliary trace segment
    /// with the specified index.
    ///
    /// The elements are drawn uniformly at random from the provided public coin.
    fn get_aux_trace_segment_random_elements<E, R>(
        &self,
        aux_segment_idx: usize,
        public_coin: &mut R,
    ) -> Result<Vec<E>, RandomCoinError>
    where
        E: FieldElement<BaseField = Self::BaseField>,
        R: RandomCoin<BaseField = Self::BaseField>,
    {
        let num_elements = self
            .trace_info()
            .layout()
            .get_aux_segment_rand_elements(aux_segment_idx);
        let mut result = Vec::with_capacity(num_elements);
        for _ in 0..num_elements {
            result.push(public_coin.draw()?);
        }
        Ok(result)
    }

    // LINEAR COMBINATION COEFFICIENTS
    // --------------------------------------------------------------------------------------------

    /// Returns coefficients needed for random linear combination during construction of constraint
    /// composition polynomial.
    fn get_constraint_composition_coefficients<E, R>(
        &self,
        public_coin: &mut R,
    ) -> Result<ConstraintCompositionCoefficients<E>, RandomCoinError>
    where
        E: FieldElement<BaseField = Self::BaseField>,
        R: RandomCoin<BaseField = Self::BaseField>,
    {
        let mut t_coefficients = Vec::new();
        for _ in 0..self.context().num_transition_constraints() {
            t_coefficients.push(public_coin.draw()?);
        }

        let mut b_coefficients = Vec::new();
        for _ in 0..self.context().num_assertions() {
            b_coefficients.push(public_coin.draw()?);
        }

        Ok(ConstraintCompositionCoefficients {
            transition: t_coefficients,
            boundary: b_coefficients,
        })
    }

    /// Returns coefficients needed for random linear combinations during construction of DEEP
    /// composition polynomial.
    fn get_deep_composition_coefficients<E, R>(
        &self,
        public_coin: &mut R,
    ) -> Result<DeepCompositionCoefficients<E>, RandomCoinError>
    where
        E: FieldElement<BaseField = Self::BaseField>,
        R: RandomCoin<BaseField = Self::BaseField>,
    {
        let mut t_coefficients = Vec::new();
        for _ in 0..self.trace_info().width() {
            t_coefficients.push(public_coin.draw()?);
        }

        let mut c_coefficients = Vec::new();
        for _ in 0..self.context().num_constraint_composition_columns() {
            c_coefficients.push(public_coin.draw()?);
        }

        Ok(DeepCompositionCoefficients {
            trace: t_coefficients,
            constraints: c_coefficients,
        })
    }
}


// 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::{
    Air, AirContext, Assertion, EvaluationFrame, ProofOptions, TraceInfo,
    TransitionConstraintDegree,
};
use crate::{AuxTraceRandElements, FieldExtension};
use crypto::{hashers::Blake3_256, DefaultRandomCoin, RandomCoin};
use math::{fields::f64::BaseElement, get_power_series, polynom, FieldElement, StarkField};
use utils::collections::{BTreeMap, Vec};

// PERIODIC COLUMNS
// ================================================================================================
#[test]
fn get_periodic_column_polys() {
    // no periodic columns
    let air = MockAir::with_periodic_columns(vec![], 16);
    let column_polys = air.get_periodic_column_polys();
    assert_eq!(0, column_polys.len());

    let col1 = vec![BaseElement::ONE, BaseElement::ZERO];
    let col2 = vec![
        BaseElement::ONE,
        BaseElement::ZERO,
        BaseElement::ONE,
        BaseElement::ONE,
    ];
    let air = MockAir::with_periodic_columns(vec![col1.clone(), col2.clone()], 16);
    let column_polys = air.get_periodic_column_polys();
    assert_eq!(2, column_polys.len());
    assert_eq!(build_periodic_column_poly(&col1), column_polys[0]);
    assert_eq!(build_periodic_column_poly(&col2), column_polys[1]);
}

#[test]
#[should_panic(expected = "number of values in a periodic column must be at least 2, but was 1")]
fn get_periodic_column_polys_num_values_too_small() {
    let col1 = vec![BaseElement::ONE];
    let air = MockAir::with_periodic_columns(vec![col1], 16);
    let column_polys = air.get_periodic_column_polys();
    assert_eq!(0, column_polys.len());
}

#[test]
#[should_panic(
    expected = "number of values in a periodic column must be a power of two, but was 3"
)]
fn get_periodic_column_polys_num_values_not_power_of_two() {
    let col1 = vec![BaseElement::ONE, BaseElement::ZERO, BaseElement::ONE];
    let air = MockAir::with_periodic_columns(vec![col1], 16);
    let column_polys = air.get_periodic_column_polys();
    assert_eq!(0, column_polys.len());
}

// TRANSITION CONSTRAINTS
// ================================================================================================

// TODO

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

#[test]
fn get_boundary_constraints() {
    // define assertions
    let values = vec![
        BaseElement::new(1),
        BaseElement::new(2),
        BaseElement::new(3),
        BaseElement::new(4),
    ];
    let assertions = vec![
        Assertion::single(0, 0, BaseElement::new(3)), // column 0, step 0 -> group 0
        Assertion::single(0, 9, BaseElement::new(5)), // column 0, step 9 -> group 1
        Assertion::single(1, 9, BaseElement::new(9)), // column 0, step 9 -> group 1
        Assertion::sequence(0, 2, 4, values.clone()), // column 0, steps 2, 6, 10, 14 -> group 2
        Assertion::sequence(1, 2, 4, values.clone()), // column 1, steps 2, 6, 10, 14 -> group 2
        Assertion::sequence(1, 0, 8, values[..2].to_vec()), // column 1, steps 0, 8 -> group 3
        Assertion::sequence(0, 3, 8, values[..2].to_vec()), // column 0, steps 3, 11 -> group 4
        Assertion::periodic(1, 3, 8, BaseElement::new(7)), // column 1, steps 3, 11 -> group 4
    ];

    // instantiate mock AIR
    let trace_length = 16;
    let air = MockAir::with_assertions(assertions, trace_length);
    let no_poly_offset = (0, BaseElement::ONE);
    let g = BaseElement::get_root_of_unity(trace_length.ilog2()); // trace domain generator

    // build coefficients for random liner combination; these will be derived for assertions
    // sorted first by stride, then by first step, and finally by column (similar to the order)
    // of assertions above
    let mut prng = build_prng();
    let mut expected_cc = BTreeMap::<usize, BaseElement>::new();
    expected_cc.insert(0, prng.draw().unwrap());
    expected_cc.insert(1, prng.draw().unwrap());
    expected_cc.insert(2, prng.draw().unwrap());
    expected_cc.insert(6, prng.draw().unwrap());
    expected_cc.insert(7, prng.draw().unwrap());
    expected_cc.insert(3, prng.draw().unwrap());
    expected_cc.insert(4, prng.draw().unwrap());
    expected_cc.insert(5, prng.draw().unwrap());

    // get boundary constraints from AIR, and sort constraint groups so that the order
    // is stable; the original order is just by degree_adjustment
    let mut prng = build_prng();
    let coefficients = (0..8)
        .map(|_| prng.draw().unwrap())
        .collect::<Vec<BaseElement>>();
    let constraints = air.get_boundary_constraints(&AuxTraceRandElements::new(), &coefficients);
    let groups = constraints.main_constraints().to_vec();

    assert_eq!(5, groups.len());

    // group 0
    let group = &groups[0];
    assert_eq!(1, group.divisor().degree());
    assert_eq!(vec![(1, g.exp(0))], group.divisor().numerator());
    assert_eq!(1, group.constraints().len());

    let constraint = &group.constraints()[0];
    assert_eq!(0, constraint.column());
    assert_eq!(vec![BaseElement::new(3)], constraint.poly());
    assert_eq!(no_poly_offset, constraint.poly_offset());
    assert_eq!(expected_cc[&0], constraint.cc().clone());

    // group 1
    let group = &groups[1];
    assert_eq!(1, group.divisor().degree());
    assert_eq!(vec![(1, g.exp(9))], group.divisor().numerator());
    assert_eq!(2, group.constraints().len());

    let constraint = &group.constraints()[0];
    assert_eq!(0, constraint.column());
    assert_eq!(vec![BaseElement::new(5)], constraint.poly());
    assert_eq!(no_poly_offset, constraint.poly_offset());
    assert_eq!(expected_cc[&1], constraint.cc().clone());

    let constraint = &group.constraints()[1];
    assert_eq!(1, constraint.column());
    assert_eq!(vec![BaseElement::new(9)], constraint.poly());
    assert_eq!(no_poly_offset, constraint.poly_offset());
    assert_eq!(expected_cc[&2], constraint.cc().clone());

    // group 2
    let group = &groups[2];
    assert_eq!(4, group.divisor().degree());
    assert_eq!(vec![(4, g.exp(4 * 2))], group.divisor().numerator());
    assert_eq!(2, group.constraints().len());

    let constraint = &group.constraints()[0];
    assert_eq!(0, constraint.column());
    assert_eq!(
        build_sequence_poly(&values, trace_length),
        constraint.poly()
    );
    assert_eq!((2, g.inv().exp(2)), constraint.poly_offset());
    assert_eq!(expected_cc[&6], constraint.cc().clone());

    let constraint = &group.constraints()[1];
    assert_eq!(1, constraint.column());
    assert_eq!(
        build_sequence_poly(&values, trace_length),
        constraint.poly()
    );
    assert_eq!((2, g.inv().exp(2)), constraint.poly_offset());
    assert_eq!(expected_cc[&7], constraint.cc().clone());

    // group 3
    let group = &groups[3];
    assert_eq!(2, group.divisor().degree());
    assert_eq!(vec![(2, g.exp(0))], group.divisor().numerator());
    assert_eq!(1, group.constraints().len());

    let constraint = &group.constraints()[0];
    assert_eq!(1, constraint.column());
    assert_eq!(
        build_sequence_poly(&values[..2], trace_length),
        constraint.poly()
    );
    assert_eq!(no_poly_offset, constraint.poly_offset());
    assert_eq!(expected_cc[&3], constraint.cc().clone());

    // group 4
    let group = &groups[4];
    assert_eq!(2, group.divisor().degree());
    assert_eq!(vec![(2, g.exp(2 * 3))], group.divisor().numerator());
    assert_eq!(2, group.constraints().len());

    let constraint = &group.constraints()[0];
    assert_eq!(0, constraint.column());
    assert_eq!(
        build_sequence_poly(&values[..2], trace_length),
        constraint.poly()
    );
    assert_eq!((3, g.inv().exp(3)), constraint.poly_offset());
    assert_eq!(expected_cc[&4], constraint.cc().clone());

    let constraint = &group.constraints()[1];
    assert_eq!(1, constraint.column());
    assert_eq!(vec![BaseElement::new(7)], constraint.poly());
    assert_eq!(no_poly_offset, constraint.poly_offset());
    assert_eq!(expected_cc[&5], constraint.cc().clone());
}

// MOCK AIR
// ================================================================================================

struct MockAir {
    context: AirContext<BaseElement>,
    assertions: Vec<Assertion<BaseElement>>,
    periodic_columns: Vec<Vec<BaseElement>>,
}

impl MockAir {
    pub fn with_periodic_columns(
        column_values: Vec<Vec<BaseElement>>,
        trace_length: usize,
    ) -> Self {
        let mut result = Self::new(
            TraceInfo::with_meta(4, trace_length, vec![1]),
            (),
            ProofOptions::new(32, 8, 0, FieldExtension::None, 4, 31),
        );
        result.periodic_columns = column_values;
        result
    }

    pub fn with_assertions(assertions: Vec<Assertion<BaseElement>>, trace_length: usize) -> Self {
        let mut result = Self::new(
            TraceInfo::with_meta(4, trace_length, vec![assertions.len() as u8]),
            (),
            ProofOptions::new(32, 8, 0, FieldExtension::None, 4, 31),
        );
        result.assertions = assertions;
        result
    }
}

impl Air for MockAir {
    type BaseField = BaseElement;
    type PublicInputs = ();

    fn new(trace_info: TraceInfo, _pub_inputs: (), _options: ProofOptions) -> Self {
        let num_assertions = trace_info.meta()[0] as usize;
        let context = build_context(trace_info.length(), trace_info.width(), num_assertions);
        MockAir {
            context,
            assertions: Vec::new(),
            periodic_columns: Vec::new(),
        }
    }

    fn context(&self) -> &AirContext<Self::BaseField> {
        &self.context
    }

    fn get_periodic_column_values(&self) -> Vec<Vec<Self::BaseField>> {
        self.periodic_columns.clone()
    }

    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
        self.assertions.clone()
    }

    fn evaluate_transition<E: FieldElement + From<Self::BaseField>>(
        &self,
        _frame: &EvaluationFrame<E>,
        _periodic_values: &[E],
        _result: &mut [E],
    ) {
    }
}

// UTILITY FUNCTIONS
// ================================================================================================

pub fn build_context<B: StarkField>(
    trace_length: usize,
    trace_width: usize,
    num_assertions: usize,
) -> AirContext<B> {
    let options = ProofOptions::new(32, 8, 0, FieldExtension::None, 4, 31);
    let t_degrees = vec![TransitionConstraintDegree::new(2)];
    let trace_info = TraceInfo::new(trace_width, trace_length);
    AirContext::new(trace_info, t_degrees, num_assertions, options)
}

pub fn build_prng() -> DefaultRandomCoin<Blake3_256<BaseElement>> {
    RandomCoin::new(&[BaseElement::ZERO; 32])
}

pub fn build_sequence_poly(values: &[BaseElement], trace_length: usize) -> Vec<BaseElement> {
    let cycle_length = trace_length / values.len();
    let domain_size = trace_length / cycle_length;
    let g = BaseElement::get_root_of_unity(domain_size.ilog2());
    let xs = get_power_series(g, domain_size);
    polynom::interpolate(&xs, values, false)
}

pub fn build_periodic_column_poly(values: &[BaseElement]) -> Vec<BaseElement> {
    let domain_size = values.len();
    let g = BaseElement::get_root_of_unity(domain_size.ilog2());
    let xs = get_power_series(g, domain_size);
    polynom::interpolate(&xs, values, false)
}


// 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 math::{StarkField, ToElements};
use utils::{
    collections::Vec, string::ToString, ByteReader, ByteWriter, Deserializable,
    DeserializationError, Serializable,
};

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

/// Number of allowed auxiliary trace segments.
const NUM_AUX_SEGMENTS: usize = 1;

// TRACE INFO
// ================================================================================================
/// Information about a specific execution trace.
///
/// Trace info consists of trace layout info, length, and optional custom metadata. Trace layout
/// specifies the number of columns for all trace segments. Currently, a trace can consist of at
/// most two segments. Metadata is just a vector of bytes and can store any values up to 64KB in
/// size.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct TraceInfo {
    layout: TraceLayout,
    length: usize,
    meta: Vec<u8>,
}

impl TraceInfo {
    /// Smallest allowed execution trace length; currently set at 8.
    pub const MIN_TRACE_LENGTH: usize = 8;
    /// Maximum number of columns in an execution trace (across all segments); currently set at 255.
    pub const MAX_TRACE_WIDTH: usize = 255;
    /// Maximum number of bytes in trace metadata; currently set at 65535.
    pub const MAX_META_LENGTH: usize = 65535;
    /// Maximum number of random elements per auxiliary trace segment; currently set to 255.
    pub const MAX_RAND_SEGMENT_ELEMENTS: usize = 255;

    // CONSTRUCTORS
    // --------------------------------------------------------------------------------------------

    /// Creates a new [TraceInfo] from the specified trace width and length.
    ///
    /// An execution trace described by this trace info is limited to a single segment.
    ///
    /// # Panics
    /// Panics if:
    /// * Trace width is zero or greater than 255.
    /// * Trace length is smaller than 8 or is not a power of two.
    pub fn new(width: usize, length: usize) -> Self {
        Self::with_meta(width, length, vec![])
    }

    /// Creates a new [TraceInfo] from the specified trace width, length, and metadata.
    ///
    /// An execution trace described by this trace info is limited to a single segment.
    ///
    /// # Panics
    /// Panics if:
    /// * Trace width is zero or greater than 255.
    /// * Trace length is smaller than 8 or is not a power of two.
    /// * Length of `meta` is greater than 65535;
    pub fn with_meta(width: usize, length: usize, meta: Vec<u8>) -> Self {
        assert!(width > 0, "trace width must be greater than 0");
        let layout = TraceLayout::new(width, [0], [0]);
        Self::new_multi_segment(layout, length, meta)
    }

    /// Creates a new [TraceInfo] from the specified trace segment widths, length, and metadata.
    ///
    /// # Panics
    /// Panics if:
    /// * The width of the first trace segment is zero.
    /// * Total width of all trace segments is greater than 255.
    /// * Trace length is smaller than 8 or is not a power of two.
    pub fn new_multi_segment(layout: TraceLayout, length: usize, meta: Vec<u8>) -> Self {
        assert!(
            length >= Self::MIN_TRACE_LENGTH,
            "trace length must be at least {}, but was {}",
            Self::MIN_TRACE_LENGTH,
            length
        );
        assert!(
            length.is_power_of_two(),
            "trace length must be a power of two, but was {length}"
        );
        assert!(
            meta.len() <= Self::MAX_META_LENGTH,
            "number of metadata bytes cannot be greater than {}, but was {}",
            Self::MAX_META_LENGTH,
            meta.len()
        );
        TraceInfo {
            layout,
            length,
            meta,
        }
    }

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

    /// Returns a description of how execution trace columns are arranged into segments.
    ///
    /// Currently, an execution trace can consist of at most two segments.
    pub fn layout(&self) -> &TraceLayout {
        &self.layout
    }

    /// Returns the total number of columns in an execution trace.
    ///
    /// This is guaranteed to be between 1 and 255.
    pub fn width(&self) -> usize {
        self.layout.main_trace_width() + self.layout().aux_trace_width()
    }

    /// Returns execution trace length.
    ///
    /// The length is guaranteed to be a power of two.
    pub fn length(&self) -> usize {
        self.length
    }

    /// Returns execution trace metadata.
    pub fn meta(&self) -> &[u8] {
        &self.meta
    }

    /// Returns true if an execution trace contains more than one segment.
    pub fn is_multi_segment(&self) -> bool {
        self.layout.num_aux_segments > 0
    }
}

// TRACE LAYOUT
// ================================================================================================

/// Layout of columns within an execution trace.
///
/// A layout describes how columns of a trace are arranged into segments. All execution traces must
/// have a non-zero main segment, and may have additional auxiliary trace segments. Currently, the
/// number of auxiliary trace segments is limited to one.
///
/// Additionally, a layout contains information on how many random elements are required to build a
/// given auxiliary trace segment. This information is used to construct
/// [AuxTraceRandElements](crate::AuxTraceRandElements) struct which is passed in as one of the
/// parameters to [Air::evaluate_aux_transition()](crate::Air::evaluate_aux_transition()) and
/// [Air::get_aux_assertions()](crate::Air::get_aux_assertions()) methods.
///
/// The number of random elements may be different from the number of columns in a given auxiliary
/// segment. For example, an auxiliary segment may contain just one column, but may require many
/// random elements.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct TraceLayout {
    main_segment_width: usize,
    aux_segment_widths: [usize; NUM_AUX_SEGMENTS],
    aux_segment_rands: [usize; NUM_AUX_SEGMENTS],
    num_aux_segments: usize,
}

impl TraceLayout {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new [TraceLayout] instantiated with the provided info.
    ///
    /// # Panics
    /// Panics if:
    /// * Width of the main trace segment is set to zero.
    /// * Sum of all segment widths exceeds 255.
    /// * A zero entry in auxiliary segment width array is followed by a non-zero entry.
    /// * Number of random elements for an auxiliary trace segment of non-zero width is set to zero.
    /// * Number of random elements for an auxiliary trace segment of zero width is set to non-zero.
    /// * Number of random elements for any auxiliary trace segment is greater than 255.
    pub fn new(
        main_width: usize,
        aux_widths: [usize; NUM_AUX_SEGMENTS],
        aux_rands: [usize; NUM_AUX_SEGMENTS],
    ) -> Self {
        // validate trace segment widths
        assert!(
            main_width > 0,
            "main trace segment must consist of at least one column"
        );
        let full_width = main_width + aux_widths.iter().sum::<usize>();
        assert!(
            full_width <= TraceInfo::MAX_TRACE_WIDTH,
            "total number of columns in the trace cannot be greater than {}, but was {}",
            TraceInfo::MAX_TRACE_WIDTH,
            full_width
        );

        // validate number of random elements required by each segment
        let mut was_zero_width = false;
        let mut num_aux_segments = 0;
        for (&width, &num_rand_elements) in aux_widths.iter().zip(aux_rands.iter()) {
            if width != 0 {
                assert!(
                    !was_zero_width,
                    "a non-empty trace segment cannot follow an empty segment"
                );
                assert!(
                    num_rand_elements > 0,
                    "number of random elements for a non-empty trace segment must be greater than zero"
                );
                num_aux_segments += 1;
            } else {
                assert!(
                    num_rand_elements == 0,
                    "number of random elements for an empty trace segment must be zero"
                );
                was_zero_width = true;
            }
            assert!(
                num_rand_elements <= TraceInfo::MAX_RAND_SEGMENT_ELEMENTS,
                "number of random elements required by a segment cannot exceed {}, but was {}",
                TraceInfo::MAX_RAND_SEGMENT_ELEMENTS,
                num_rand_elements
            );
        }

        Self {
            main_segment_width: main_width,
            aux_segment_widths: aux_widths,
            aux_segment_rands: aux_rands,
            num_aux_segments,
        }
    }

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

    /// Returns the number of columns in the main segment of an execution trace.
    ///
    /// This is guaranteed to be between 1 and 255.
    pub fn main_trace_width(&self) -> usize {
        self.main_segment_width
    }

    /// Returns the number of columns in all auxiliary segments of an execution trace.
    pub fn aux_trace_width(&self) -> usize {
        self.aux_segment_widths.iter().sum()
    }

    /// Returns the total number of segments in an execution trace.
    pub fn num_segments(&self) -> usize {
        self.num_aux_segments + 1
    }

    /// Returns the number of auxiliary trace segments in an execution trace.
    pub fn num_aux_segments(&self) -> usize {
        self.num_aux_segments
    }

    /// Returns the number of columns in the auxiliary trace segment at the specified index.
    pub fn get_aux_segment_width(&self, segment_idx: usize) -> usize {
        // TODO: panic if segment_idx is not within num_aux_segments
        self.aux_segment_widths[segment_idx]
    }

    /// Returns the number of random elements required by the auxiliary trace segment at the
    /// specified index.
    pub fn get_aux_segment_rand_elements(&self, segment_idx: usize) -> usize {
        // TODO: panic if segment_idx is not within num_aux_segments
        self.aux_segment_rands[segment_idx]
    }
}

impl<E: StarkField> ToElements<E> for TraceLayout {
    fn to_elements(&self) -> Vec<E> {
        let mut result = Vec::new();

        // main segment width, number of auxiliary segments, and parameters of the first auxiliary
        // segment (if present) go into the first field element; we assume that each parameter can
        // be encoded in 8 bits (which is enforced by the constructor)
        let mut buf = self.main_segment_width as u32;
        buf = (buf << 8) | self.num_aux_segments as u32;
        if self.num_aux_segments == 1 {
            buf = (buf << 8) | self.aux_segment_widths[0] as u32;
            buf = (buf << 8) | self.aux_segment_rands[0] as u32;
        }
        result.push(E::from(buf));

        // parameters of all subsequent auxiliary segments go into additional elements
        for i in 1..self.num_aux_segments {
            buf = self.aux_segment_widths[i] as u32;
            buf = (buf << 8) | self.aux_segment_rands[i] as u32;
            result.push(E::from(buf));
        }

        result
    }
}

impl Serializable for TraceLayout {
    /// Serializes `self` and writes the resulting bytes into the `target`.
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        target.write_u8(self.main_segment_width as u8);
        for &w in self.aux_segment_widths.iter() {
            debug_assert!(
                w <= u8::MAX as usize,
                "aux segment width does not fit into u8 value"
            );
            target.write_u8(w as u8);
        }
        for &rc in self.aux_segment_rands.iter() {
            debug_assert!(
                rc <= u8::MAX as usize,
                "aux segment random element count does not fit into u8 value"
            );
            target.write_u8(rc as u8);
        }
    }
}

impl Deserializable for TraceLayout {
    /// Reads [TraceLayout] from the specified `source` and returns the result.
    ///
    /// # Errors
    /// Returns an error of a valid [TraceLayout] struct could not be read from the specified
    /// `source`.
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let main_width = source.read_u8()? as usize;
        if main_width == 0 {
            return Err(DeserializationError::InvalidValue(
                "main trace segment width must be greater than zero".to_string(),
            ));
        }

        // read and validate auxiliary trace segment widths
        let mut was_zero_width = false;
        let mut aux_widths = [0; NUM_AUX_SEGMENTS];
        for width in aux_widths.iter_mut() {
            *width = source.read_u8()? as usize;
            if *width != 0 {
                if was_zero_width {
                    return Err(DeserializationError::InvalidValue(
                        "a non-empty trace segment cannot follow an empty segment".to_string(),
                    ));
                }
            } else {
                was_zero_width = true;
            }
        }

        let full_trace_width = main_width + aux_widths.iter().sum::<usize>();
        if full_trace_width >= TraceInfo::MAX_TRACE_WIDTH {
            return Err(DeserializationError::InvalidValue(format!(
                "full trace width cannot be greater than {}, but was {}",
                TraceInfo::MAX_TRACE_WIDTH,
                full_trace_width
            )));
        }

        // read and validate number of random elements for each auxiliary trace segment
        let mut aux_rands = [0; NUM_AUX_SEGMENTS];
        for (num_rand_elements, &width) in aux_rands.iter_mut().zip(aux_widths.iter()) {
            *num_rand_elements = source.read_u8()? as usize;
            if width == 0 && *num_rand_elements != 0 {
                return Err(DeserializationError::InvalidValue(
                    "an empty trace segment cannot require random elements".to_string(),
                ));
            } else if width != 0 && *num_rand_elements == 0 {
                return Err(DeserializationError::InvalidValue(
                    "a non-empty trace segment must require at least one random element"
                        .to_string(),
                ));
            } else if *num_rand_elements > TraceInfo::MAX_RAND_SEGMENT_ELEMENTS {
                return Err(DeserializationError::InvalidValue(format!(
                    "number of random elements required by a segment cannot exceed {}, but was {}",
                    TraceInfo::MAX_RAND_SEGMENT_ELEMENTS,
                    *num_rand_elements
                )));
            }
        }

        Ok(TraceLayout::new(main_width, aux_widths, aux_rands))
    }
}

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

#[cfg(test)]
mod tests {
    use super::{ToElements, TraceLayout};
    use math::fields::f64::BaseElement;

    #[test]
    fn trace_layout_to_elements() {
        // --- test trace with only main segment ------------------------------
        let main_width = 20;
        let num_aux_segments = 0;

        let expected = u32::from_le_bytes([num_aux_segments, main_width as u8, 0, 0]);
        let expected = vec![BaseElement::from(expected)];

        let layout = TraceLayout::new(main_width, [0], [0]);
        assert_eq!(expected, layout.to_elements());

        // --- test trace with one auxiliary segment --------------------------
        let main_width = 20;
        let num_aux_segments = 1;
        let aux_width = 9;
        let aux_rands = 12;

        let expected = u32::from_le_bytes([aux_rands, aux_width, num_aux_segments, main_width]);
        let expected = vec![BaseElement::from(expected)];

        let layout = TraceLayout::new(
            main_width as usize,
            [aux_width as usize],
            [aux_rands as usize],
        );
        assert_eq!(expected, layout.to_elements());
    }
}