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

//! Basic polynomial operations.
//!
//! This module provides a set of function for basic polynomial operations, including:
//! - Polynomial evaluation using Horner method.
//! - Polynomial interpolation using Lagrange method.
//! - Polynomial addition, subtraction, multiplication, and division.
//! - Synthetic polynomial division for efficient division by polynomials of the form
//!   `x`^`a` - `b`.
//!
//! In the context of this module any slice of field elements is considered to be a polynomial
//! in reverse coefficient form. A few examples:
//!
//! ```
//! # use winter_math::{fields::{f128::BaseElement}, FieldElement};
//! // p(x) = 2 * x + 1
//! let p = vec![BaseElement::new(1), BaseElement::new(2)];
//!
//! // p(x) = 4 * x^2 + 3
//! let p = [BaseElement::new(3), BaseElement::ZERO, BaseElement::new(4)];
//! ```

use crate::{field::FieldElement, utils::batch_inversion};
use core::mem;
use utils::{collections::Vec, group_vector_elements};

#[cfg(test)]
mod tests;

// POLYNOMIAL EVALUATION
// ================================================================================================

/// Evaluates a polynomial at a single point and returns the result.
///
/// Evaluates polynomial `p` at coordinate `x` using
/// [Horner's method](https://en.wikipedia.org/wiki/Horner%27s_method).
///
/// # Examples
/// ```
/// # use winter_math::polynom::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// // define polynomial: f(x) = 3 * x^2 + 2 * x + 1
/// let p = (1u32..4).map(BaseElement::from).collect::<Vec<_>>();
///
/// // evaluate the polynomial at point 4
/// let x = BaseElement::new(4);
/// assert_eq!(BaseElement::new(57), eval(&p, x));
/// ```
pub fn eval<B, E>(p: &[B], x: E) -> E
where
    B: FieldElement,
    E: FieldElement + From<B>,
{
    // Horner evaluation
    p.iter()
        .rev()
        .fold(E::ZERO, |acc, &coeff| acc * x + E::from(coeff))
}

/// Evaluates a polynomial at multiple points and returns a vector of results.
///
/// Evaluates polynomial `p` at all coordinates in `xs` slice by repeatedly invoking
/// `polynom::eval()` function.
///
/// # Examples
/// ```
/// # use winter_math::polynom::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// // define polynomial: f(x) = 3 * x^2 + 2 * x + 1
/// let p = (1_u32..4).map(BaseElement::from).collect::<Vec<_>>();
/// let xs = (3_u32..6).map(BaseElement::from).collect::<Vec<_>>();
///
/// let expected = xs.iter().map(|x| eval(&p, *x)).collect::<Vec<_>>();
/// assert_eq!(expected, eval_many(&p, &xs));
/// ```
pub fn eval_many<B, E>(p: &[B], xs: &[E]) -> Vec<E>
where
    B: FieldElement,
    E: FieldElement + From<B>,
{
    xs.iter().map(|x| eval(p, *x)).collect()
}

// POLYNOMIAL INTERPOLATION
// ================================================================================================

/// Returns a polynomial in coefficient form interpolated from a set of X and Y coordinates.
///
/// Uses [Lagrange interpolation](https://en.wikipedia.org/wiki/Lagrange_polynomial) to build a
/// polynomial from X and Y coordinates. If `remove_leading_zeros = true`, all leading coefficients
/// which are ZEROs will be truncated; otherwise, the length of result will be equal to the number
/// of X coordinates.
///
/// # Panics
/// Panics if number of X and Y coordinates is not the same.
///
/// # Example
/// ```
/// # use winter_math::polynom::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// # use rand_utils::rand_vector;
/// let xs: Vec<BaseElement> = rand_vector(16);
/// let ys: Vec<BaseElement> = rand_vector(16);
///
/// let p = interpolate(&xs, &ys, false);
/// assert_eq!(ys, eval_many(&p, &xs));
/// ```
pub fn interpolate<E>(xs: &[E], ys: &[E], remove_leading_zeros: bool) -> Vec<E>
where
    E: FieldElement,
{
    debug_assert!(
        xs.len() == ys.len(),
        "number of X and Y coordinates must be the same"
    );

    let roots = get_zero_roots(xs);
    let numerators: Vec<Vec<E>> = xs.iter().map(|&x| syn_div(&roots, 1, x)).collect();

    let denominators: Vec<E> = numerators
        .iter()
        .zip(xs)
        .map(|(e, &x)| eval(e, x))
        .collect();
    let denominators = batch_inversion(&denominators);

    let mut result = E::zeroed_vector(xs.len());
    for i in 0..xs.len() {
        let y_slice = ys[i] * denominators[i];
        for (j, res) in result.iter_mut().enumerate() {
            *res += numerators[i][j] * y_slice;
        }
    }

    if remove_leading_zeros {
        crate::polynom::remove_leading_zeros(&result)
    } else {
        result
    }
}

/// Returns a vector of polynomials interpolated from the provided X and Y coordinate batches.
///
/// Uses [Lagrange interpolation](https://en.wikipedia.org/wiki/Lagrange_polynomial) to build a
/// vector of polynomial from X and Y coordinate batches (one polynomial per batch).
///
/// When the number of batches is larger, this function is significantly faster than using
/// `polynom::interpolate()` function individually for each batch of coordinates. The speed-up
/// is primarily due to computing all inversions as a single batch inversion across all
/// coordinate batches.
///
/// # Panics
/// Panics if the number of X coordinate batches and Y coordinate batches is not the same.
///
/// # Examples
/// ```
/// # use core::convert::TryInto;
/// # use winter_math::polynom::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// # use rand_utils::rand_array;
/// let x_batches: Vec<[BaseElement; 8]> = vec![
///     rand_array(),
///     rand_array(),
/// ];
/// let y_batches: Vec<[BaseElement; 8]> = vec![
///     rand_array(),
///     rand_array(),
/// ];
///
/// let polys = interpolate_batch(&x_batches, &y_batches);
/// for ((p, xs), ys) in polys.iter().zip(x_batches).zip(y_batches) {
///     assert_eq!(ys.to_vec(), eval_many(p, &xs));
/// }
/// ```
pub fn interpolate_batch<E, const N: usize>(xs: &[[E; N]], ys: &[[E; N]]) -> Vec<[E; N]>
where
    E: FieldElement,
{
    debug_assert!(
        xs.len() == ys.len(),
        "number of X coordinate batches and Y coordinate batches must be the same"
    );

    let n = xs.len();
    let mut equations = group_vector_elements(E::zeroed_vector(n * N * N));
    let mut inverses = E::zeroed_vector(n * N);

    // TODO: converting this to an array results in about 5% speed-up, but unfortunately, complex
    // generic constraints are not yet supported: https://github.com/rust-lang/rust/issues/76560
    let mut roots = vec![E::ZERO; N + 1];

    for (i, xs) in xs.iter().enumerate() {
        fill_zero_roots(xs, &mut roots);
        for (j, &x) in xs.iter().enumerate() {
            let equation = &mut equations[i * N + j];
            // optimized synthetic division for this context
            equation[N - 1] = roots[N];
            for k in (0..N - 1).rev() {
                equation[k] = roots[k + 1] + equation[k + 1] * x;
            }
            inverses[i * N + j] = eval(equation, x);
        }
    }
    let equations = group_vector_elements::<[E; N], N>(equations);
    let inverses = group_vector_elements::<E, N>(batch_inversion(&inverses));

    let mut result = group_vector_elements(E::zeroed_vector(n * N));
    for (i, poly) in result.iter_mut().enumerate() {
        for j in 0..N {
            let inv_y = ys[i][j] * inverses[i][j];
            for (res_coeff, &eq_coeff) in poly.iter_mut().zip(equations[i][j].iter()) {
                *res_coeff += eq_coeff * inv_y;
            }
        }
    }

    result
}

// POLYNOMIAL MATH OPERATIONS
// ================================================================================================

/// Returns a polynomial resulting from adding two polynomials together.
///
/// Polynomials `a` and `b` are expected to be in the coefficient form, and the returned
/// polynomial will be in the coefficient form as well. The length of the returned vector
/// will be max(a.len(), b.len()).
///
/// # Examples
/// ```
/// # use winter_math::polynom::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// // p1(x) = 4 * x^2 + 3 * x + 2
/// let p1 = (2_u32..5).map(BaseElement::from).collect::<Vec<_>>();
/// // p2(x) = 2 * x + 1
/// let p2 = (1_u32..3).map(BaseElement::from).collect::<Vec<_>>();
///
/// // expected result = 4 * x^2 + 5 * x + 3
/// let expected = vec![
///     BaseElement::new(3),
///     BaseElement::new(5),
///     BaseElement::new(4),
/// ];
/// assert_eq!(expected, add(&p1, &p2));
/// ```
pub fn add<E>(a: &[E], b: &[E]) -> Vec<E>
where
    E: FieldElement,
{
    let result_len = core::cmp::max(a.len(), b.len());
    let mut result = Vec::with_capacity(result_len);
    for i in 0..result_len {
        let c1 = if i < a.len() { a[i] } else { E::ZERO };
        let c2 = if i < b.len() { b[i] } else { E::ZERO };
        result.push(c1 + c2);
    }
    result
}

/// Returns a polynomial resulting from subtracting one polynomial from another.
///
/// Specifically, subtracts polynomial `b` from polynomial `a` and returns the result. Both
/// polynomials are expected to be in the coefficient form, and the returned polynomial will
/// be in the coefficient form as well.  The length of the returned vector will be
/// max(a.len(), b.len()).
///
/// # Examples
/// ```
/// # use winter_math::polynom::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// // p1(x) = 4 * x^2 + 3 * x + 2
/// let p1 = (2_u32..5).map(BaseElement::from).collect::<Vec<_>>();
/// // p2(x) = 2 * x + 1
/// let p2 = (1_u32..3).map(BaseElement::from).collect::<Vec<_>>();
///
/// // expected result = 4 * x^2 + x + 1
/// let expected = vec![
///     BaseElement::new(1),
///     BaseElement::new(1),
///     BaseElement::new(4),
/// ];
/// assert_eq!(expected, sub(&p1, &p2));
/// ```
pub fn sub<E>(a: &[E], b: &[E]) -> Vec<E>
where
    E: FieldElement,
{
    let result_len = core::cmp::max(a.len(), b.len());
    let mut result = Vec::with_capacity(result_len);
    for i in 0..result_len {
        let c1 = if i < a.len() { a[i] } else { E::ZERO };
        let c2 = if i < b.len() { b[i] } else { E::ZERO };
        result.push(c1 - c2);
    }
    result
}

/// Returns a polynomial resulting from multiplying two polynomials together.
///
/// Polynomials `a` and `b` are expected to be in the coefficient form, and the returned
/// polynomial will be in the coefficient form as well. The length of the returned vector
/// will be a.len() + b.len() - 1.
///
/// # Examples
/// ```
/// # use winter_math::polynom::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// // p1(x) = x + 1
/// let p1 = [BaseElement::ONE, BaseElement::ONE];
/// // p2(x) = x^2 + 2
/// let p2 = [BaseElement::new(2), BaseElement::ZERO, BaseElement::ONE];
///
/// // expected result = x^3 + x^2 + 2 * x + 2
/// let expected = vec![
///     BaseElement::new(2),
///     BaseElement::new(2),
///     BaseElement::new(1),
///     BaseElement::new(1),
/// ];
/// assert_eq!(expected, mul(&p1, &p2));
/// ```
pub fn mul<E>(a: &[E], b: &[E]) -> Vec<E>
where
    E: FieldElement,
{
    let result_len = a.len() + b.len() - 1;
    let mut result = E::zeroed_vector(result_len);
    for i in 0..a.len() {
        for j in 0..b.len() {
            let s = a[i] * b[j];
            result[i + j] += s;
        }
    }
    result
}

/// Returns a polynomial resulting from multiplying a given polynomial by a scalar value.
///
/// Specifically, multiplies every coefficient of polynomial `p` by constant `k` and returns
/// the resulting vector.
///
/// # Examples
/// ```
/// # use winter_math::polynom::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// let p = [
///     BaseElement::new(1),
///     BaseElement::new(2),
///     BaseElement::new(3),
/// ];
/// let k = BaseElement::new(2);
///
/// let expected = vec![
///     BaseElement::new(2),
///     BaseElement::new(4),
///     BaseElement::new(6),
/// ];
/// assert_eq!(expected, mul_by_scalar(&p, k));
/// ```
pub fn mul_by_scalar<E>(p: &[E], k: E) -> Vec<E>
where
    E: FieldElement,
{
    let mut result = Vec::with_capacity(p.len());
    for coeff in p {
        result.push(*coeff * k);
    }
    result
}

/// Returns a polynomial resulting from dividing one polynomial by another.
///
/// Specifically, divides polynomial `a` by polynomial `b` and returns the result. If the
/// polynomials don't divide evenly, the remainder is ignored. Both polynomials are expected to
/// be in the coefficient form, and the returned polynomial will be in the coefficient form as
/// well. The length of the returned vector will be a.len() - b.len() + 1.
///
/// # Panics
/// Panics if:
/// * Polynomial `b` is empty.
/// * Degree of polynomial `b` is zero and the constant coefficient is ZERO.
/// * The degree of polynomial `b` is greater than the degree of polynomial `a`.
///
/// # Examples
/// ```
/// # use winter_math::polynom::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// // p1(x) = x^3 + x^2 + 2 * x + 2
/// let p1 = [
///     BaseElement::new(2),
///     BaseElement::new(2),
///     BaseElement::new(1),
///     BaseElement::new(1),
/// ];
/// // p2(x) = x^2 + 2
/// let p2 = [BaseElement::new(2), BaseElement::ZERO, BaseElement::ONE];
///
/// // expected result = x + 1
/// let expected = vec![BaseElement::ONE, BaseElement::ONE];
/// assert_eq!(expected, div(&p1, &p2));
/// ```
pub fn div<E>(a: &[E], b: &[E]) -> Vec<E>
where
    E: FieldElement,
{
    let mut apos = degree_of(a);
    let mut a = a.to_vec();

    let bpos = degree_of(b);
    assert!(apos >= bpos, "cannot divide by polynomial of higher degree");
    if bpos == 0 {
        assert!(!b.is_empty(), "cannot divide by empty polynomial");
        assert!(b[0] != E::ZERO, "cannot divide polynomial by zero");
    }

    let mut result = E::zeroed_vector(apos - bpos + 1);
    for i in (0..result.len()).rev() {
        let quot = a[apos] / b[bpos];
        result[i] = quot;
        for j in (0..bpos).rev() {
            a[i + j] -= b[j] * quot;
        }
        apos = apos.wrapping_sub(1);
    }

    result
}

/// Returns a polynomial resulting from dividing a polynomial by a polynomial of special form.
///
/// Specifically, divides polynomial `p` by polynomial (x^`a` - `b`) using
/// [synthetic division](https://en.wikipedia.org/wiki/Synthetic_division) method; if the
/// polynomials don't divide evenly, the remainder is ignored. Polynomial `p` is expected
/// to be in the coefficient form, and the result will be in the coefficient form as well.
/// The length of the resulting polynomial will be equal to `p.len()`.
///
/// This function is significantly faster than the generic `polynom::div()` function.
///
/// # Panics
/// Panics if:
/// * `a` is zero;
/// * `b` is zero;
/// * `p.len()` is smaller than or equal to `a`.
///
/// # Examples
/// ```
/// # use winter_math::polynom::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// // p(x) = x^3 + x^2 + 2 * x + 2
/// let p = [
///     BaseElement::new(2),
///     BaseElement::new(2),
///     BaseElement::new(1),
///     BaseElement::new(1),
/// ];
///
/// // expected result = x^2 + 2
/// let expected = vec![
///     BaseElement::new(2),
///     BaseElement::ZERO,
///     BaseElement::new(1),
///     BaseElement::ZERO,
/// ];
///
/// // divide by x + 1
/// assert_eq!(expected, syn_div(&p, 1, -BaseElement::ONE));
/// ```
pub fn syn_div<E>(p: &[E], a: usize, b: E) -> Vec<E>
where
    E: FieldElement,
{
    let mut result = p.to_vec();
    syn_div_in_place(&mut result, a, b);
    result
}

/// Divides a polynomial by a polynomial of special form and saves the result into the original
/// polynomial.
///
/// Specifically, divides polynomial `p` by polynomial (x^`a` - `b`) using
/// [synthetic division](https://en.wikipedia.org/wiki/Synthetic_division) method and saves the
/// result into `p`. If the polynomials don't divide evenly, the remainder is ignored. Polynomial
/// `p` is expected to be in the coefficient form, and the result will be in coefficient form as
/// well.
///
/// This function is significantly faster than the generic `polynom::div()` function, and as
/// compared to `polynom::syn_div()` function, this function does not allocate any additional
/// memory.
///
/// # Panics
/// Panics if:
/// * `a` is zero;
/// * `b` is zero;
/// * `p.len()` is smaller than or equal to `a`.
///
/// # Examples
/// ```
/// # use winter_math::polynom::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// // p(x) = x^3 + x^2 + 2 * x + 2
/// let mut p = [
///     BaseElement::new(2),
///     BaseElement::new(2),
///     BaseElement::new(1),
///     BaseElement::new(1),
/// ];
///
/// // divide by x + 1
/// syn_div_in_place(&mut p, 1, -BaseElement::ONE);
///
/// // expected result = x^2 + 2
/// let expected = [
///     BaseElement::new(2),
///     BaseElement::ZERO,
///     BaseElement::new(1),
///     BaseElement::ZERO,
/// ];
///
/// assert_eq!(expected, p);
pub fn syn_div_in_place<E>(p: &mut [E], a: usize, b: E)
where
    E: FieldElement,
{
    assert!(a != 0, "divisor degree cannot be zero");
    assert!(b != E::ZERO, "constant cannot be zero");
    assert!(
        p.len() > a,
        "divisor degree cannot be greater than dividend size"
    );

    if a == 1 {
        // if we are dividing by (x - `b`), we can use a single variable to keep track
        // of the remainder; this way, we can avoid shifting the values in the slice later
        let mut c = E::ZERO;
        for coeff in p.iter_mut().rev() {
            *coeff += b * c;
            mem::swap(coeff, &mut c);
        }
    } else {
        // if we are dividing by a polynomial of higher power, we need to keep track of the
        // full remainder. we do that in place, but then need to shift the values at the end
        // to discard the remainder
        let degree_offset = p.len() - a;
        if b == E::ONE {
            // if `b` is 1, no need to multiply by `b` in every iteration of the loop
            for i in (0..degree_offset).rev() {
                p[i] += p[i + a];
            }
        } else {
            for i in (0..degree_offset).rev() {
                p[i] += p[i + a] * b;
            }
        }
        // discard the remainder
        p.copy_within(a.., 0);
        p[degree_offset..].fill(E::ZERO);
    }
}

// DEGREE INFERENCE
// ================================================================================================

/// Returns the degree of the provided polynomial.
///
/// If the size of the provided slice is much larger than the degree of the polynomial (i.e.,
/// a large number of leading coefficients is ZERO), this operation can be quite inefficient.
///
/// # Examples
/// ```
/// # use winter_math::polynom::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// assert_eq!(0, degree_of::<BaseElement>(&[]));
/// assert_eq!(0, degree_of(&[BaseElement::ONE]));
/// assert_eq!(1, degree_of(&[BaseElement::ONE, BaseElement::new(2)]));
/// assert_eq!(1, degree_of(&[BaseElement::ONE, BaseElement::new(2), BaseElement::ZERO]));
/// assert_eq!(2, degree_of(&[BaseElement::ONE, BaseElement::new(2), BaseElement::new(3)]));
/// assert_eq!(
///     2,
///     degree_of(&[
///         BaseElement::ONE,
///         BaseElement::new(2),
///         BaseElement::new(3),
///         BaseElement::ZERO
///     ])
/// );
/// ```
pub fn degree_of<E>(poly: &[E]) -> usize
where
    E: FieldElement,
{
    for i in (0..poly.len()).rev() {
        if poly[i] != E::ZERO {
            return i;
        }
    }
    0
}

/// Returns a polynomial with all leading ZERO coefficients removed.
///
/// # Examples
/// ```
/// # use winter_math::polynom::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// let a = vec![1u128, 2, 3, 4, 5, 6, 0, 0]
///     .into_iter()
///     .map(BaseElement::new)
///     .collect::<Vec<_>>();
/// let b = remove_leading_zeros(&a);
/// assert_eq!(6, b.len());
/// assert_eq!(a[..6], b);
///
/// let a = vec![0u128, 0, 0, 0]
///     .into_iter()
///     .map(BaseElement::new)
///     .collect::<Vec<_>>();
/// let b = remove_leading_zeros(&a);
/// assert_eq!(0, b.len());
/// ```
pub fn remove_leading_zeros<E>(values: &[E]) -> Vec<E>
where
    E: FieldElement,
{
    for i in (0..values.len()).rev() {
        if values[i] != E::ZERO {
            return values[..(i + 1)].to_vec();
        }
    }
    vec![]
}

// HELPER FUNCTIONS
// ================================================================================================
fn get_zero_roots<E: FieldElement>(xs: &[E]) -> Vec<E> {
    let mut result = unsafe { utils::uninit_vector(xs.len() + 1) };
    fill_zero_roots(xs, &mut result);
    result
}

fn fill_zero_roots<E: FieldElement>(xs: &[E], result: &mut [E]) {
    let mut n = result.len();
    n -= 1;
    result[n] = E::ONE;

    for i in 0..xs.len() {
        n -= 1;
        result[n] = E::ZERO;
        #[allow(clippy::assign_op_pattern)]
        for j in n..xs.len() {
            result[j] = result[j] - result[j + 1] * xs[i];
        }
    }
}


// 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::remove_leading_zeros;
use crate::{
    field::{f128::BaseElement, FieldElement, StarkField},
    utils::get_power_series,
};
use utils::collections::Vec;

#[test]
fn eval() {
    let x = BaseElement::from(11269864713250585702u128);
    let poly: [BaseElement; 4] = [
        BaseElement::from(384863712573444386u128),
        BaseElement::from(7682273369345308472u128),
        BaseElement::from(13294661765012277990u128),
        BaseElement::from(16234810094004944758u128),
    ];

    assert_eq!(BaseElement::ZERO, super::eval(&poly[..0], x));

    // constant
    assert_eq!(poly[0], super::eval(&poly[..1], x));

    // degree 1
    assert_eq!(poly[0] + poly[1] * x, super::eval(&poly[..2], x));

    // degree 2
    let x2 = x.exp(2);
    assert_eq!(
        poly[0] + poly[1] * x + poly[2] * x2,
        super::eval(&poly[..3], x)
    );

    // degree 3
    let x3 = x.exp(3);
    assert_eq!(
        poly[0] + poly[1] * x + poly[2] * x2 + poly[3] * x3,
        super::eval(&poly, x)
    );
}

#[test]
fn add() {
    let poly1: [BaseElement; 3] = [
        BaseElement::from(384863712573444386u128),
        BaseElement::from(7682273369345308472u128),
        BaseElement::from(13294661765012277990u128),
    ];
    let poly2: [BaseElement; 3] = [
        BaseElement::from(9918505539874556741u128),
        BaseElement::from(16401861429499852246u128),
        BaseElement::from(12181445947541805654u128),
    ];

    // same degree
    let pr = vec![
        poly1[0] + poly2[0],
        poly1[1] + poly2[1],
        poly1[2] + poly2[2],
    ];
    assert_eq!(pr, super::add(&poly1, &poly2));

    // poly1 is lower degree
    let pr = vec![poly1[0] + poly2[0], poly1[1] + poly2[1], poly2[2]];
    assert_eq!(pr, super::add(&poly1[..2], &poly2));

    // poly2 is lower degree
    let pr = vec![poly1[0] + poly2[0], poly1[1] + poly2[1], poly1[2]];
    assert_eq!(pr, super::add(&poly1, &poly2[..2]));
}

#[test]
fn sub() {
    let poly1: [BaseElement; 3] = [
        BaseElement::from(384863712573444386u128),
        BaseElement::from(7682273369345308472u128),
        BaseElement::from(13294661765012277990u128),
    ];
    let poly2: [BaseElement; 3] = [
        BaseElement::from(9918505539874556741u128),
        BaseElement::from(16401861429499852246u128),
        BaseElement::from(12181445947541805654u128),
    ];

    // same degree
    let pr = vec![
        poly1[0] - poly2[0],
        poly1[1] - poly2[1],
        poly1[2] - poly2[2],
    ];
    assert_eq!(pr, super::sub(&poly1, &poly2));

    // poly1 is lower degree
    let pr = vec![poly1[0] - poly2[0], poly1[1] - poly2[1], -poly2[2]];
    assert_eq!(pr, super::sub(&poly1[..2], &poly2));

    // poly2 is lower degree
    let pr = vec![poly1[0] - poly2[0], poly1[1] - poly2[1], poly1[2]];
    assert_eq!(pr, super::sub(&poly1, &poly2[..2]));
}

#[test]
fn mul() {
    let poly1: [BaseElement; 3] = [
        BaseElement::from(384863712573444386u128),
        BaseElement::from(7682273369345308472u128),
        BaseElement::from(13294661765012277990u128),
    ];
    let poly2: [BaseElement; 3] = [
        BaseElement::from(9918505539874556741u128),
        BaseElement::from(16401861429499852246u128),
        BaseElement::from(12181445947541805654u128),
    ];

    // same degree
    let pr = vec![
        poly1[0] * poly2[0],
        poly1[0] * poly2[1] + poly2[0] * poly1[1],
        poly1[1] * poly2[1] + poly1[2] * poly2[0] + poly2[2] * poly1[0],
        poly1[2] * poly2[1] + poly2[2] * poly1[1],
        poly1[2] * poly2[2],
    ];
    assert_eq!(pr, super::mul(&poly1, &poly2));

    // poly1 is lower degree
    let pr = vec![
        poly1[0] * poly2[0],
        poly1[0] * poly2[1] + poly2[0] * poly1[1],
        poly1[0] * poly2[2] + poly2[1] * poly1[1],
        poly1[1] * poly2[2],
    ];
    assert_eq!(pr, super::mul(&poly1[..2], &poly2));

    // poly2 is lower degree
    let pr = vec![
        poly1[0] * poly2[0],
        poly1[0] * poly2[1] + poly2[0] * poly1[1],
        poly1[2] * poly2[0] + poly2[1] * poly1[1],
        poly1[2] * poly2[1],
    ];
    assert_eq!(pr, super::mul(&poly1, &poly2[..2]));
}

#[test]
fn div() {
    let poly1 = vec![
        BaseElement::from(384863712573444386u128),
        BaseElement::from(7682273369345308472u128),
        BaseElement::from(13294661765012277990u128),
    ];
    let poly2 = vec![
        BaseElement::from(9918505539874556741u128),
        BaseElement::from(16401861429499852246u128),
        BaseElement::from(12181445947541805654u128),
    ];

    // divide degree 4 by degree 2
    let poly3 = super::mul(&poly1, &poly2);
    assert_eq!(poly1, super::div(&poly3, &poly2));

    // divide degree 3 by degree 2
    let poly3 = super::mul(&poly1[..2], &poly2);
    assert_eq!(poly1[..2].to_vec(), super::div(&poly3, &poly2));

    // divide degree 3 by degree 3
    let poly3 = super::mul_by_scalar(&poly1, BaseElement::from(11269864713250585702u128));
    assert_eq!(
        vec![BaseElement::from(11269864713250585702u128)],
        super::div(&poly3, &poly1)
    );
}

#[test]
fn syn_div() {
    // ----- division by degree 1 polynomial ------------------------------------------------------

    // poly = (x + 2) * (x + 3)
    let poly = super::mul(
        &[BaseElement::from(2u8), BaseElement::ONE],
        &[BaseElement::from(3u8), BaseElement::ONE],
    );

    // divide by (x + 3), this divides evenly
    let result = super::syn_div(&poly, 1, -BaseElement::from(3u8));
    let expected = vec![BaseElement::from(2u8), BaseElement::ONE];
    assert_eq!(expected, remove_leading_zeros(&result));

    // poly = x^3 - 12x^2 - 42
    let poly = [
        -BaseElement::from(42u8),
        BaseElement::ZERO,
        -BaseElement::from(12u8),
        BaseElement::ONE,
    ];

    // divide by (x - 3), this does not divide evenly, but the remainder is ignored
    let result = super::syn_div(&poly, 1, BaseElement::from(3u8));
    let expected = vec![
        -BaseElement::from(27u8),
        -BaseElement::from(9u8),
        BaseElement::ONE,
    ];
    assert_eq!(expected, remove_leading_zeros(&result));

    // ----- division by high-degree polynomial ---------------------------------------------------

    // evaluations of a polynomial which evaluates to 0 at steps: 0, 4, 8, 12
    let ys: Vec<BaseElement> = vec![0u8, 1, 2, 3, 0, 5, 6, 7, 0, 9, 10, 11, 0, 13, 14, 15]
        .into_iter()
        .map(BaseElement::from)
        .collect();

    // build the domain
    let root = BaseElement::get_root_of_unity(ys.len().ilog2());
    let domain = get_power_series(root, ys.len());

    // build the polynomial
    let poly = super::interpolate(&domain, &ys, false);

    // build the divisor polynomial: (x^4 - 1)
    let z_poly = vec![
        -BaseElement::ONE,
        BaseElement::ZERO,
        BaseElement::ZERO,
        BaseElement::ZERO,
        BaseElement::ONE,
    ];

    let result = super::syn_div(&poly, 4, BaseElement::ONE);
    assert_eq!(poly, remove_leading_zeros(&super::mul(&result, &z_poly)));

    // ----- division by high-degree polynomial with non-unary constant ---------------------------

    // evaluations of a polynomial which evaluates to 0 at steps: 1, 5, 9, 13
    let ys: Vec<BaseElement> = vec![18u8, 0, 2, 3, 4, 0, 6, 7, 8, 0, 10, 11, 12, 0, 14, 15]
        .into_iter()
        .map(BaseElement::from)
        .collect();

    // build the polynomial
    let poly = super::interpolate(&domain, &ys, false);

    // build the divisor polynomial: (x^4 - g^4)
    let z_poly = vec![
        -root.exp(4),
        BaseElement::ZERO,
        BaseElement::ZERO,
        BaseElement::ZERO,
        BaseElement::ONE,
    ];

    let result = super::syn_div(&poly, 4, root.exp(4));
    assert_eq!(poly, remove_leading_zeros(&super::mul(&result, &z_poly)));
}


// 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::{field::FieldElement, ExtensionOf};
use utils::{batch_iter_mut, collections::Vec, iter_mut, uninit_vector};

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

// MATH FUNCTIONS
// ================================================================================================

/// Returns a vector containing successive powers of a given base.
///
/// More precisely, for base `b`, generates a vector with values [1, b, b^2, b^3, ..., b^(n-1)].
///
/// When `concurrent` feature is enabled, series generation is done concurrently in multiple
/// threads.
///
/// # Examples
/// ```
/// # use winter_math::get_power_series;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// let n = 2048;
/// let b = BaseElement::from(3u8);
///
/// let expected = (0..n)
///     .map(|p| b.exp((p as u64).into()))
///     .collect::<Vec<_>>();
///
/// let actual = get_power_series(b, n);
/// assert_eq!(expected, actual);
/// ```
pub fn get_power_series<E>(b: E, n: usize) -> Vec<E>
where
    E: FieldElement,
{
    let mut result = unsafe { uninit_vector(n) };
    batch_iter_mut!(&mut result, 1024, |batch: &mut [E], batch_offset: usize| {
        let start = b.exp((batch_offset as u64).into());
        fill_power_series(batch, b, start);
    });
    result
}

/// Returns a vector containing successive powers of a given base offset by the specified value.
///
/// More precisely, for base `b` and offset `s`, generates a vector with values
/// [s, s * b, s * b^2, s * b^3, ..., s * b^(n-1)].
///
/// When `concurrent` feature is enabled, series generation is done concurrently in multiple
/// threads.
///
/// # Examples
/// ```
/// # use winter_math::get_power_series_with_offset;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// let n = 2048;
/// let b = BaseElement::from(3u8);
/// let s = BaseElement::from(7u8);
///
/// let expected = (0..n)
///     .map(|p| s * b.exp((p as u64).into()))
///     .collect::<Vec<_>>();
///
/// let actual = get_power_series_with_offset(b, s, n);
/// assert_eq!(expected, actual);
/// ```
pub fn get_power_series_with_offset<E>(b: E, s: E, n: usize) -> Vec<E>
where
    E: FieldElement,
{
    let mut result = unsafe { uninit_vector(n) };
    batch_iter_mut!(&mut result, 1024, |batch: &mut [E], batch_offset: usize| {
        let start = s * b.exp((batch_offset as u64).into());
        fill_power_series(batch, b, start);
    });
    result
}

/// Computes element-wise sum of the provided vectors, and stores the result in the first vector.
///
/// When `concurrent` feature is enabled, the summation is performed concurrently in multiple
/// threads.
///
/// # Panics
/// Panics if lengths of `a` and `b` vectors are not the same.
///
/// # Examples
/// ```
/// # use winter_math::add_in_place;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// # use rand_utils::rand_vector;
/// let a: Vec<BaseElement> = rand_vector(2048);
/// let b: Vec<BaseElement> = rand_vector(2048);
///
/// let mut c = a.clone();
/// add_in_place(&mut c, &b);
///
/// for ((a, b), c) in a.into_iter().zip(b).zip(c) {
///     assert_eq!(a + b, c);
/// }
/// ```
pub fn add_in_place<E>(a: &mut [E], b: &[E])
where
    E: FieldElement,
{
    assert!(
        a.len() == b.len(),
        "number of values must be the same for both operands"
    );
    iter_mut!(a).zip(b).for_each(|(a, &b)| *a += b);
}

/// Multiplies a sequence of values by a scalar and accumulates the results.
///
/// More precisely, computes `a[i]` + `b[i]` * `c` for all `i` and saves result into `a[i]`.
///
/// When `concurrent` feature is enabled, the computation is performed concurrently in multiple
/// threads.
///
/// # Panics
/// Panics if lengths of `a` and `b` slices are not the same.
///
/// # Examples
/// ```
/// # use winter_math::mul_acc;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// # use rand_utils::rand_vector;
/// let a: Vec<BaseElement> = rand_vector(2048);
/// let b: Vec<BaseElement> = rand_vector(2048);
/// let c = BaseElement::new(12345);
///
/// let mut d = a.clone();
/// mul_acc(&mut d, &b, c);
///
/// for ((a, b), d) in a.into_iter().zip(b).zip(d) {
///     assert_eq!(a + b * c, d);
/// }
/// ```
pub fn mul_acc<F, E>(a: &mut [E], b: &[F], c: E)
where
    F: FieldElement,
    E: FieldElement<BaseField = F::BaseField> + ExtensionOf<F>,
{
    assert!(
        a.len() == b.len(),
        "number of values must be the same for both slices"
    );
    iter_mut!(a).zip(b).for_each(|(a, &b)| *a += c.mul_base(b));
}

/// Computes a multiplicative inverse of a sequence of elements using batch inversion method.
///
/// Any ZEROs in the provided sequence are ignored.
///
/// When `concurrent` feature is enabled, the inversion is performed concurrently in multiple
/// threads.
///
/// This function is significantly faster than inverting elements one-by-one because it
/// essentially transforms `n` inversions into `3 * n` multiplications + 1 inversion.
///
/// # Examples
/// ```
/// # use winter_math::batch_inversion;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
/// # use rand_utils::rand_vector;
/// let a: Vec<BaseElement> = rand_vector(2048);
/// let b = batch_inversion(&a);
///
/// for (&a, &b) in a.iter().zip(b.iter()) {
///     assert_eq!(a.inv(), b);
/// }
/// ```
pub fn batch_inversion<E>(values: &[E]) -> Vec<E>
where
    E: FieldElement,
{
    let mut result: Vec<E> = unsafe { uninit_vector(values.len()) };
    batch_iter_mut!(&mut result, 1024, |batch: &mut [E], batch_offset: usize| {
        let start = batch_offset;
        let end = start + batch.len();
        serial_batch_inversion(&values[start..end], batch);
    });
    result
}

/// Returns base 2 logarithm of `n`, where `n` is a power of two.
///
/// # Panics
/// Panics if `n` is not a power of two.
///
/// # Examples
/// ```
/// # use winter_math::log2;
/// assert_eq!(log2(1), 0);
/// assert_eq!(log2(16), 4);
/// assert_eq!(log2(1 << 20), 20);
/// assert_eq!(log2(2usize.pow(20)), 20);
/// ```
pub fn log2(n: usize) -> u32 {
    assert!(n.is_power_of_two(), "n must be a power of two");
    n.trailing_zeros()
}

// HELPER FUNCTIONS
// ------------------------------------------------------------------------------------------------

#[inline(always)]
fn fill_power_series<E: FieldElement>(result: &mut [E], base: E, start: E) {
    result[0] = start;
    for i in 1..result.len() {
        result[i] = result[i - 1] * base;
    }
}

fn serial_batch_inversion<E: FieldElement>(values: &[E], result: &mut [E]) {
    let mut last = E::ONE;
    for (result, &value) in result.iter_mut().zip(values.iter()) {
        *result = last;
        if value != E::ZERO {
            last *= value;
        }
    }

    last = last.inv();

    for i in (0..values.len()).rev() {
        if values[i] == E::ZERO {
            result[i] = E::ZERO;
        } else {
            result[i] *= last;
            last *= values[i];
        }
    }
}


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

use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use math::{fft, fields::f64::BaseElement, StarkField};
use rand_utils::rand_vector;
use std::time::Duration;
use winter_prover::{ColMatrix, RowMatrix, StarkDomain};

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

const SIZE: usize = 524_288;
const BLOWUP_FACTOR: [usize; 3] = [2, 4, 8];
const NUM_POLYS: [usize; 3] = [32, 64, 96];

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

    for &num_poly in NUM_POLYS.iter() {
        let columns: Vec<Vec<BaseElement>> = (0..num_poly).map(|_| rand_vector(SIZE)).collect();
        let column_matrix = ColMatrix::new(columns);
        for &blowup_factor in BLOWUP_FACTOR.iter() {
            let params = BenchmarkParams {
                num_poly,
                blowup_factor,
            };
            group.bench_function(BenchmarkId::new(SIZE.to_string(), params), |bench| {
                bench.iter_with_large_drop(|| {
                    let twiddles = fft::get_twiddles::<BaseElement>(SIZE);
                    let stark_domain =
                        StarkDomain::from_twiddles(twiddles, blowup_factor, BaseElement::GENERATOR);
                    column_matrix.evaluate_columns_over(&stark_domain)
                });
            });
        }
    }
    group.finish();
}

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

    for &num_poly in NUM_POLYS.iter() {
        let columns: Vec<Vec<BaseElement>> = (0..num_poly).map(|_| rand_vector(SIZE)).collect();
        let column_matrix = ColMatrix::new(columns);
        for &blowup_factor in BLOWUP_FACTOR.iter() {
            let params = BenchmarkParams {
                num_poly,
                blowup_factor,
            };
            group.bench_function(BenchmarkId::new(SIZE.to_string(), params), |bench| {
                bench.iter_with_large_drop(|| {
                    RowMatrix::evaluate_polys::<8>(&column_matrix, blowup_factor);
                });
            });
        }
    }
    group.finish();
}

/// Benchmark parameters.
struct BenchmarkParams {
    num_poly: usize,
    blowup_factor: usize,
}

impl std::fmt::Display for BenchmarkParams {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "num_poly: {}, blowup_factor: {}",
            self.num_poly, self.blowup_factor
        )
    }
}

criterion_group!(matrix_group, evaluate_columns, evaluate_matrix);
criterion_main!(matrix_group);


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

use air::{
    proof::{Commitments, Context, OodFrame, Queries, StarkProof},
    Air, ConstraintCompositionCoefficients, DeepCompositionCoefficients,
};
use core::marker::PhantomData;
use crypto::{ElementHasher, RandomCoin};
use fri::{self, FriProof};
use math::{FieldElement, ToElements};
use utils::collections::Vec;

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

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

pub struct ProverChannel<'a, A, E, H, R>
where
    A: Air,
    E: FieldElement<BaseField = A::BaseField>,
    H: ElementHasher<BaseField = A::BaseField>,
    R: RandomCoin<BaseField = E::BaseField, Hasher = H>,
{
    air: &'a A,
    public_coin: R,
    context: Context,
    commitments: Commitments,
    ood_frame: OodFrame,
    pow_nonce: u64,
    _field_element: PhantomData<E>,
}

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

impl<'a, A, E, H, R> ProverChannel<'a, A, E, H, R>
where
    A: Air,
    E: FieldElement<BaseField = A::BaseField>,
    H: ElementHasher<BaseField = A::BaseField>,
    R: RandomCoin<BaseField = A::BaseField, Hasher = H>,
{
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Creates a new prover channel for the specified `air` and public inputs.
    pub fn new(air: &'a A, mut pub_inputs_elements: Vec<A::BaseField>) -> Self {
        let context = Context::new::<A::BaseField>(air.trace_info(), air.options().clone());

        // build a seed for the public coin; the initial seed is a hash of the proof context and
        // the public inputs, but as the protocol progresses, the coin will be reseeded with the
        // info sent to the verifier
        let mut coin_seed_elements = context.to_elements();
        coin_seed_elements.append(&mut pub_inputs_elements);

        ProverChannel {
            air,
            public_coin: RandomCoin::new(&coin_seed_elements),
            context,
            commitments: Commitments::default(),
            ood_frame: OodFrame::default(),
            pow_nonce: 0,
            _field_element: PhantomData,
        }
    }

    // COMMITMENT METHODS
    // --------------------------------------------------------------------------------------------

    /// Commits the prover the extended execution trace.
    pub fn commit_trace(&mut self, trace_root: H::Digest) {
        self.commitments.add::<H>(&trace_root);
        self.public_coin.reseed(trace_root);
    }

    /// Commits the prover to the evaluations of the constraint composition polynomial.
    pub fn commit_constraints(&mut self, constraint_root: H::Digest) {
        self.commitments.add::<H>(&constraint_root);
        self.public_coin.reseed(constraint_root);
    }

    /// Saves the evaluations of trace polynomials over the out-of-domain evaluation frame. This
    /// also reseeds the public coin with the hashes of the evaluation frame states.
    pub fn send_ood_trace_states(&mut self, trace_states: &[Vec<E>]) {
        let result = self.ood_frame.set_trace_states(trace_states);
        self.public_coin.reseed(H::hash_elements(&result));
    }

    /// Saves the evaluations of constraint composition polynomial columns at the out-of-domain
    /// point. This also reseeds the public coin wit the hash of the evaluations.
    pub fn send_ood_constraint_evaluations(&mut self, evaluations: &[E]) {
        self.ood_frame.set_constraint_evaluations(evaluations);
        self.public_coin.reseed(H::hash_elements(evaluations));
    }

    // PUBLIC COIN METHODS
    // --------------------------------------------------------------------------------------------

    /// Returns a set of random elements required for constructing an auxiliary trace segment with
    /// the specified index.
    ///
    /// The elements are drawn from the public coin uniformly at random.
    pub fn get_aux_trace_segment_rand_elements(&mut self, aux_segment_idx: usize) -> Vec<E> {
        self.air
            .get_aux_trace_segment_random_elements(aux_segment_idx, &mut self.public_coin)
            .expect("failed to draw random elements for an auxiliary trace segment")
    }

    /// Returns a set of coefficients for constructing a constraint composition polynomial.
    ///
    /// The coefficients are drawn from the public coin uniformly at random.
    pub fn get_constraint_composition_coeffs(&mut self) -> ConstraintCompositionCoefficients<E> {
        self.air
            .get_constraint_composition_coefficients(&mut self.public_coin)
            .expect("failed to draw composition coefficients")
    }

    /// Returns an out-of-domain point drawn uniformly at random from the public coin.
    pub fn get_ood_point(&mut self) -> E {
        self.public_coin.draw().expect("failed to draw OOD point")
    }

    /// Returns a set of coefficients for constructing a DEEP composition polynomial.
    ///
    /// The coefficients are drawn from the public coin uniformly at random.
    pub fn get_deep_composition_coeffs(&mut self) -> DeepCompositionCoefficients<E> {
        self.air
            .get_deep_composition_coefficients(&mut self.public_coin)
            .expect("failed to draw DEEP composition coefficients")
    }

    /// Returns a set of positions in the LDE domain against which the evaluations of trace and
    /// constraint composition polynomials should be queried.
    ///
    /// The positions are drawn from the public coin uniformly at random.
    pub fn get_query_positions(&mut self) -> Vec<usize> {
        let num_queries = self.context.options().num_queries();
        let lde_domain_size = self.context.lde_domain_size();
        self.public_coin
            .draw_integers(num_queries, lde_domain_size)
            .expect("failed to draw query position")
    }

    /// Determines a nonce, which when hashed with the current seed of the public coin results
    /// in a new seed with the number of leading zeros equal to the grinding_factor specified
    /// in the proof options.
    pub fn grind_query_seed(&mut self) {
        let grinding_factor = self.context.options().grinding_factor();

        #[cfg(not(feature = "concurrent"))]
        let nonce = (1..u64::MAX)
            .find(|&nonce| self.public_coin.check_leading_zeros(nonce) >= grinding_factor)
            .expect("nonce not found");

        #[cfg(feature = "concurrent")]
        let nonce = (1..u64::MAX)
            .into_par_iter()
            .find_any(|&nonce| self.public_coin.check_leading_zeros(nonce) >= grinding_factor)
            .expect("nonce not found");

        self.pow_nonce = nonce;
        self.public_coin.reseed_with_int(nonce);
    }

    // PROOF BUILDER
    // --------------------------------------------------------------------------------------------
    /// Builds a proof from the previously committed values as well as values passed into
    /// this method.
    pub fn build_proof(
        self,
        trace_queries: Vec<Queries>,
        constraint_queries: Queries,
        fri_proof: FriProof,
    ) -> StarkProof {
        StarkProof {
            context: self.context,
            commitments: self.commitments,
            ood_frame: self.ood_frame,
            trace_queries,
            constraint_queries,
            fri_proof,
            pow_nonce: self.pow_nonce,
        }
    }
}

// FRI PROVER CHANNEL IMPLEMENTATION
// ================================================================================================

impl<'a, A, E, H, R> fri::ProverChannel<E> for ProverChannel<'a, A, E, H, R>
where
    A: Air,
    E: FieldElement<BaseField = A::BaseField>,
    H: ElementHasher<BaseField = A::BaseField>,
    R: RandomCoin<BaseField = A::BaseField, Hasher = H>,
{
    type Hasher = H;

    /// Commits the prover to a FRI layer.
    fn commit_fri_layer(&mut self, layer_root: H::Digest) {
        self.commitments.add::<H>(&layer_root);
        self.public_coin.reseed(layer_root);
    }

    /// Returns a new alpha drawn from the public coin.
    fn draw_fri_alpha(&mut self) -> E {
        self.public_coin.draw().expect("failed to draw FRI alpha")
    }
}


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

use air::Air;
use math::{fft, get_power_series, StarkField};
use utils::collections::Vec;

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

pub struct StarkDomain<B: StarkField> {
    /// Twiddles which can be used to evaluate polynomials in the trace domain. Length of this
    /// vector is half the length of the trace domain size.
    trace_twiddles: Vec<B>,

    /// [g^i for i in (0..ce_domain_size)] where g is the constraint evaluation domain generator.
    ce_domain: Vec<B>,

    /// LDE domain size / constraint evaluation domain size
    ce_to_lde_blowup: usize,

    /// A mask which can be used to compute (x % ce_domain_size) via binary AND. This takes
    /// advantage of the fact that ce_domain_size is a power of two. The mask is then simply
    /// ce_domain_size - 1.
    ce_domain_mod_mask: usize,

    /// Offset of the low-degree extension domain.
    domain_offset: B,
}

// STARK DOMAIN IMPLEMENTATION
// ================================================================================================

impl<B: StarkField> StarkDomain<B> {
    /// Returns a new STARK domain initialized with the provided `context`.
    pub fn new<A: Air<BaseField = B>>(air: &A) -> Self {
        let trace_twiddles = fft::get_twiddles(air.trace_length());

        // build constraint evaluation domain
        let domain_gen = B::get_root_of_unity(air.ce_domain_size().ilog2());
        let ce_domain = get_power_series(domain_gen, air.ce_domain_size());

        StarkDomain {
            trace_twiddles,
            ce_domain,
            ce_to_lde_blowup: air.lde_domain_size() / air.ce_domain_size(),
            ce_domain_mod_mask: air.ce_domain_size() - 1,
            domain_offset: air.domain_offset(),
        }
    }

    /// Returns a new STARK domain initialized with the provided custom inputs.
    pub fn from_twiddles(trace_twiddles: Vec<B>, blowup_factor: usize, domain_offset: B) -> Self {
        // both `trace_twiddles` length and `blowup_factor` must be a power of two.
        assert!(
            trace_twiddles.len().is_power_of_two(),
            "the length of trace twiddles must be a power of 2"
        );
        assert!(
            blowup_factor.is_power_of_two(),
            "blowup factor must be a power of 2"
        );

        let ce_domain_size = trace_twiddles.len() * blowup_factor * 2;
        let domain_gen = B::get_root_of_unity(ce_domain_size.ilog2());
        let ce_domain = get_power_series(domain_gen, ce_domain_size);

        StarkDomain {
            trace_twiddles,
            ce_domain,
            ce_to_lde_blowup: 1,
            ce_domain_mod_mask: ce_domain_size - 1,
            domain_offset,
        }
    }

    // EXECUTION TRACE
    // --------------------------------------------------------------------------------------------

    /// Returns length of the execution trace for this computation.
    pub fn trace_length(&self) -> usize {
        &self.trace_twiddles.len() * 2
    }

    /// Returns twiddles which can be used to evaluate trace polynomials.
    pub fn trace_twiddles(&self) -> &[B] {
        &self.trace_twiddles
    }

    /// Returns blowup factor from trace to constraint evaluation domain.
    pub fn trace_to_ce_blowup(&self) -> usize {
        self.ce_domain_size() / self.trace_length()
    }

    /// Returns blowup factor from trace to LDE domain.
    pub fn trace_to_lde_blowup(&self) -> usize {
        self.lde_domain_size() / self.trace_length()
    }

    // CONSTRAINT EVALUATION DOMAIN
    // --------------------------------------------------------------------------------------------

    /// Returns the size of the constraint evaluation domain for this computation.
    #[inline(always)]
    pub fn ce_domain_size(&self) -> usize {
        self.ce_domain.len()
    }

    /// Returns the generator of constraint evaluation domain.
    pub fn ce_domain_generator(&self) -> B {
        B::get_root_of_unity(self.ce_domain_size().ilog2())
    }

    /// Returns blowup factor from constraint evaluation to LDE domain.
    pub fn ce_to_lde_blowup(&self) -> usize {
        self.ce_to_lde_blowup
    }

    /// Returns s * g^step where g is the constraint evaluation domain generator and s is the
    /// domain offset.
    #[inline(always)]
    pub fn get_ce_x_at(&self, step: usize) -> B {
        self.ce_domain[step] * self.domain_offset
    }

    /// Returns (s * g^step)^power where g is the constraint evaluation domain generator and s is
    /// the domain offset.
    ///
    /// The computation is performed without doing exponentiations. offset_exp is assumed to be
    /// s^power which is pre-computed elsewhere.
    #[inline(always)]
    pub fn get_ce_x_power_at(&self, step: usize, power: u64, offset_exp: B) -> B {
        debug_assert_eq!(offset_exp, self.offset().exp(power.into()));
        // this computes (step * power) % ce_domain_size. even though both step and power could be
        // 64-bit values, we are not concerned about overflow here because we are modding by a
        // power of two. this is also the reason why we can do & ce_domain_mod_mask instead of
        // performing the actual modulus operation.
        let index = step.wrapping_mul(power as usize) & self.ce_domain_mod_mask;
        self.ce_domain[index] * offset_exp
    }

    // LOW-DEGREE EXTENSION DOMAIN
    // --------------------------------------------------------------------------------------------

    /// Returns the size of the low-degree extension domain.
    pub fn lde_domain_size(&self) -> usize {
        self.ce_domain_size() * self.ce_to_lde_blowup()
    }

    /// Returns LDE domain offset.
    pub fn offset(&self) -> B {
        self.domain_offset
    }
}


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

//! Contains common error types for prover and verifier.

use core::fmt;

// PROVER ERROR
// ================================================================================================
/// Represents an error returned by the prover during an execution of the protocol.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProverError {
    /// This error occurs when a transition constraint evaluated over a specific execution trace
    /// does not evaluate to zero at any of the steps.
    UnsatisfiedTransitionConstraintError(usize),
    /// This error occurs when polynomials built from the columns of a constraint evaluation
    /// table do not all have the same degree.
    MismatchedConstraintPolynomialDegree(usize, usize),
    /// This error occurs when the base field specified by the AIR does not support field extension
    /// of degree specified by proof options.
    UnsupportedFieldExtension(usize),
}

impl fmt::Display for ProverError {
    #[rustfmt::skip]
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::UnsatisfiedTransitionConstraintError(step) => {
                write!(f, "a transition constraint was not satisfied at step {step}")
            }
            Self::MismatchedConstraintPolynomialDegree(expected, actual) => {
                write!(f, "the constraint polynomial's components do not all have the same degree; expected {expected}, but was {actual}")
            }
            Self::UnsupportedFieldExtension(degree) => {
                write!(f, "field extension of degree {degree} is not supported for the specified base field")
            }
        }
    }
}


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

//! This crate contains Winterfell STARK prover.
//!
//! This prover can be used to generate proofs of computational integrity using the
//! [STARK](https://eprint.iacr.org/2018/046) (Scalable Transparent ARguments of Knowledge)
//! protocol.
//!
//! When the crate is compiled with `concurrent` feature enabled, proof generation will be
//! performed in multiple threads (usually, as many threads as there are logical cores on the
//! machine). The number of threads can be configured via `RAYON_NUM_THREADS` environment
//! variable.
//!
//! # Usage
//! To generate a proof that a computation was executed correctly, you'll need to do the
//! following:
//!
//! 1. Define an *algebraic intermediate representation* (AIR) for your computation. This can
//!    be done by implementing [Air] trait.
//! 2. Define an execution trace for your computation. This can be done by implementing [Trace]
//!    trait. Alternatively, you can use [TraceTable] struct which already implements [Trace]
//!    trait in cases when this generic implementation works for your use case.
//! 3. Execute your computation and record its execution trace.
//! 4. Define your prover by implementing [Prover] trait. Then execute [Prover::prove()] function
//!    passing the trace generated in the previous step into it as a parameter. The function will
//!    return a instance of [StarkProof].
//!
//! This [StarkProof] can be serialized and sent to a STARK verifier for verification. The size
//! of proof depends on the specifics of a given computation, but for most computations it should
//! be in the range between 15 KB (for very small computations) and 300 KB (for very large
//! computations).
//!
//! Proof generation time is also highly dependent on the specifics of a given computation, but
//! also depends on the capabilities of the machine used to generate the proofs (i.e. on number
//! of CPU cores and memory bandwidth).

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

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

pub use air::{
    proof::StarkProof, Air, AirContext, Assertion, AuxTraceRandElements, BoundaryConstraint,
    BoundaryConstraintGroup, ConstraintCompositionCoefficients, ConstraintDivisor,
    DeepCompositionCoefficients, EvaluationFrame, FieldExtension, ProofOptions, TraceInfo,
    TraceLayout, TransitionConstraintDegree,
};
pub use utils::{
    iterators, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
    SliceReader,
};

use fri::FriProver;
use utils::collections::Vec;

pub use math;
use math::{
    fft::infer_degree,
    fields::{CubeExtension, QuadExtension},
    ExtensibleField, FieldElement, StarkField, ToElements,
};

pub use crypto;
use crypto::{ElementHasher, MerkleTree, RandomCoin};

#[cfg(feature = "std")]
use log::debug;
#[cfg(feature = "std")]
use std::time::Instant;

mod domain;
pub use domain::StarkDomain;

pub mod matrix;
pub use matrix::{ColMatrix, RowMatrix};

mod constraints;
use constraints::ConstraintEvaluator;
pub use constraints::{CompositionPoly, ConstraintCommitment};

mod composer;
use composer::DeepCompositionPoly;

mod trace;
pub use trace::{Trace, TraceTable, TraceTableFragment};
use trace::{TraceCommitment, TraceLde, TracePolyTable};

mod channel;
use channel::ProverChannel;

mod errors;
pub use errors::ProverError;

#[cfg(test)]
pub mod tests;

// PROVER
// ================================================================================================

// this segment width seems to give the best performance for small fields (i.e., 64 bits)
const DEFAULT_SEGMENT_WIDTH: usize = 8;

/// Defines a STARK prover for a computation.
///
/// A STARK prover can be used to generate STARK proofs. The prover contains definitions of a
/// computation's AIR (specified via [Air](Prover::Air) associated type), execution trace
/// (specified via [Trace](Prover::Trace) associated type) and hash function to be used (specified
/// via [HashFn](Prover::HashFn) associated type), and exposes [prove()](Prover::prove) method which can
/// be used to build STARK proofs for provided execution traces.
///
/// Thus, once a prover is defined and instantiated, generating a STARK proof consists of two
/// steps:
/// 1. Build an execution trace for a specific instance of the computation.
/// 2. Invoke [Prover::prove()] method generate a proof using the trace from the previous step
///    as a witness.
///
/// The generated proof is built using protocol parameters defined by the [ProofOptions] struct
/// return from [Prover::options] method.
pub trait Prover {
    /// Base field for the computation described by this prover.
    type BaseField: StarkField + ExtensibleField<2> + ExtensibleField<3>;

    /// Algebraic intermediate representation (AIR) for the computation described by this prover.
    type Air: Air<BaseField = Self::BaseField>;

    /// Execution trace of the computation described by this prover.
    type Trace: Trace<BaseField = Self::BaseField>;

    /// Hash function to be used.
    type HashFn: ElementHasher<BaseField = Self::BaseField>;

    /// PRNG to be used for generating random field elements.
    type RandomCoin: RandomCoin<BaseField = Self::BaseField, Hasher = Self::HashFn>;

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

    /// Returns a set of public inputs for an instance of the computation defined by the provided
    /// trace.
    ///
    /// Public inputs need to be shared with the verifier in order for them to verify a proof.
    fn get_pub_inputs(&self, trace: &Self::Trace) -> <<Self as Prover>::Air as Air>::PublicInputs;

    /// Returns [ProofOptions] which this prover uses to generate STARK proofs.
    ///
    /// Proof options defines basic protocol parameters such as: number of queries, blowup factor,
    /// grinding factor etc. These properties directly inform such metrics as proof generation time,
    /// proof size, and proof security level.
    fn options(&self) -> &ProofOptions;

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

    /// Returns a STARK proof attesting to a correct execution of a computation defined by the
    /// provided trace.
    ///
    /// The returned [StarkProof] attests that the specified `trace` is a valid execution trace of
    /// the computation described by [Self::Air](Prover::Air) and generated using some set of
    /// secret and public inputs. Public inputs must match the value returned from
    /// [Self::get_pub_inputs()](Prover::get_pub_inputs) for the provided trace.
    #[rustfmt::skip]
    fn prove(&self, trace: Self::Trace) -> Result<StarkProof, ProverError> {
        // figure out which version of the generic proof generation procedure to run. this is a sort
        // of static dispatch for selecting two generic parameter: extension field and hash function.
        match self.options().field_extension() {
            FieldExtension::None => self.generate_proof::<Self::BaseField>(trace),
            FieldExtension::Quadratic => {
                if !<QuadExtension<Self::BaseField>>::is_supported() {
                    return Err(ProverError::UnsupportedFieldExtension(2));
                }
                self.generate_proof::<QuadExtension<Self::BaseField>>(trace)
            }
            FieldExtension::Cubic => {
                if !<CubeExtension<Self::BaseField>>::is_supported() {
                    return Err(ProverError::UnsupportedFieldExtension(3));
                }
                self.generate_proof::<CubeExtension<Self::BaseField>>(trace)
            }
        }
    }

    // HELPER METHODS
    // --------------------------------------------------------------------------------------------

    /// Performs the actual proof generation procedure, generating the proof that the provided
    /// execution `trace` is valid against this prover's AIR.
    /// TODO: make this function un-callable externally?
    #[doc(hidden)]
    fn generate_proof<E>(&self, mut trace: Self::Trace) -> Result<StarkProof, ProverError>
    where
        E: FieldElement<BaseField = Self::BaseField>,
    {
        // 0 ----- instantiate AIR and prover channel ---------------------------------------------

        // serialize public inputs; these will be included in the seed for the public coin
        let pub_inputs = self.get_pub_inputs(&trace);
        let pub_inputs_elements = pub_inputs.to_elements();

        // create an instance of AIR for the provided parameters. this takes a generic description
        // of the computation (provided via AIR type), and creates a description of a specific
        // execution of the computation for the provided public inputs.
        let air = Self::Air::new(trace.get_info(), pub_inputs, self.options().clone());

        // create a channel which is used to simulate interaction between the prover and the
        // verifier; the channel will be used to commit to values and to draw randomness that
        // should come from the verifier.
        let mut channel = ProverChannel::<Self::Air, E, Self::HashFn, Self::RandomCoin>::new(
            &air,
            pub_inputs_elements,
        );

        // 1 ----- Commit to the execution trace --------------------------------------------------

        // build computation domain; this is used later for polynomial evaluations
        #[cfg(feature = "std")]
        let now = Instant::now();
        let domain = StarkDomain::new(&air);
        #[cfg(feature = "std")]
        debug!(
            "Built domain of 2^{} elements in {} ms",
            domain.lde_domain_size().ilog2(),
            now.elapsed().as_millis()
        );

        // extend the main execution trace and build a Merkle tree from the extended trace
        let (main_trace_lde, main_trace_tree, main_trace_polys) =
            self.build_trace_commitment::<Self::BaseField>(trace.main_segment(), &domain);

        // commit to the LDE of the main trace by writing the root of its Merkle tree into
        // the channel
        channel.commit_trace(*main_trace_tree.root());

        // initialize trace commitment and trace polynomial table structs with the main trace
        // data; for multi-segment traces these structs will be used as accumulators of all
        // trace segments
        let mut trace_commitment = TraceCommitment::new(
            main_trace_lde,
            main_trace_tree,
            domain.trace_to_lde_blowup(),
        );
        let mut trace_polys = TracePolyTable::new(main_trace_polys);

        // build auxiliary trace segments (if any), and append the resulting segments to trace
        // commitment and trace polynomial table structs
        let mut aux_trace_segments = Vec::new();
        let mut aux_trace_rand_elements = AuxTraceRandElements::new();
        for i in 0..trace.layout().num_aux_segments() {
            #[cfg(feature = "std")]
            let now = Instant::now();

            // draw a set of random elements required to build an auxiliary trace segment
            let rand_elements = channel.get_aux_trace_segment_rand_elements(i);

            // build the trace segment
            let aux_segment = trace
                .build_aux_segment(&aux_trace_segments, &rand_elements)
                .expect("failed build auxiliary trace segment");
            #[cfg(feature = "std")]
            debug!(
                "Built auxiliary trace segment of {} columns and 2^{} steps in {} ms",
                aux_segment.num_cols(),
                aux_segment.num_rows().ilog2(),
                now.elapsed().as_millis()
            );

            // extend the auxiliary trace segment and build a Merkle tree from the extended trace
            let (aux_segment_lde, aux_segment_tree, aux_segment_polys) =
                self.build_trace_commitment::<E>(&aux_segment, &domain);

            // commit to the LDE of the extended auxiliary trace segment  by writing the root of
            // its Merkle tree into the channel
            channel.commit_trace(*aux_segment_tree.root());

            // append the segment to the trace commitment and trace polynomial table structs
            trace_commitment.add_segment(aux_segment_lde, aux_segment_tree);
            trace_polys.add_aux_segment(aux_segment_polys);
            aux_trace_rand_elements.add_segment_elements(rand_elements);
            aux_trace_segments.push(aux_segment);
        }

        // make sure the specified trace (including auxiliary segments) is valid against the AIR.
        // This checks validity of both, assertions and state transitions. We do this in debug
        // mode only because this is a very expensive operation.
        #[cfg(debug_assertions)]
        trace.validate(&air, &aux_trace_segments, &aux_trace_rand_elements);

        // 2 ----- evaluate constraints -----------------------------------------------------------
        // evaluate constraints specified by the AIR over the constraint evaluation domain, and
        // compute random linear combinations of these evaluations using coefficients drawn from
        // the channel; this step evaluates only constraint numerators, thus, only constraints with
        // identical denominators are merged together. the results are saved into a constraint
        // evaluation table where each column contains merged evaluations of constraints with
        // identical denominators.
        #[cfg(feature = "std")]
        let now = Instant::now();
        let constraint_coeffs = channel.get_constraint_composition_coeffs();
        let evaluator = ConstraintEvaluator::new(&air, aux_trace_rand_elements, constraint_coeffs);
        let constraint_evaluations = evaluator.evaluate(trace_commitment.trace_table(), &domain);
        #[cfg(feature = "std")]
        debug!(
            "Evaluated constraints over domain of 2^{} elements in {} ms",
            constraint_evaluations.num_rows().ilog2(),
            now.elapsed().as_millis()
        );

        // 3 ----- commit to constraint evaluations -----------------------------------------------

        // first, build constraint composition polynomial from the constraint evaluation table:
        // - divide all constraint evaluation columns by their respective divisors
        // - combine them into a single column of evaluations,
        // - interpolate the column into a polynomial in coefficient form
        // - "break" the polynomial into a set of column polynomials each of degree equal to
        //   trace_length - 1
        #[cfg(feature = "std")]
        let now = Instant::now();
        let composition_poly =
            constraint_evaluations.into_poly(air.context().num_constraint_composition_columns())?;
        #[cfg(feature = "std")]
        debug!(
            "Converted constraint evaluations into {} composition polynomial columns of degree {} in {} ms",
            composition_poly.num_columns(),
            composition_poly.column_degree(),
            now.elapsed().as_millis()
        );

        // then, build a commitment to the evaluations of the composition polynomial columns
        let constraint_commitment =
            self.build_constraint_commitment::<E>(&composition_poly, &domain);

        // then, commit to the evaluations of constraints by writing the root of the constraint
        // Merkle tree into the channel
        channel.commit_constraints(constraint_commitment.root());

        // 4 ----- build DEEP composition polynomial ----------------------------------------------
        #[cfg(feature = "std")]
        let now = Instant::now();

        // draw an out-of-domain point z. Depending on the type of E, the point is drawn either
        // from the base field or from an extension field defined by E.
        //
        // The purpose of sampling from the extension field here (instead of the base field) is to
        // increase security. Soundness is limited by the size of the field that the random point
        // is drawn from, and we can potentially save on performance by only drawing this point
        // from an extension field, rather than increasing the size of the field overall.
        let z = channel.get_ood_point();

        // evaluate trace and constraint polynomials at the OOD point z, and send the results to
        // the verifier. the trace polynomials are actually evaluated over two points: z and z * g,
        // where g is the generator of the trace domain.
        let ood_trace_states = trace_polys.get_ood_frame(z);
        channel.send_ood_trace_states(&ood_trace_states);

        let ood_evaluations = composition_poly.evaluate_at(z);
        channel.send_ood_constraint_evaluations(&ood_evaluations);

        // draw random coefficients to use during DEEP polynomial composition, and use them to
        // initialize the DEEP composition polynomial
        let deep_coefficients = channel.get_deep_composition_coeffs();
        let mut deep_composition_poly = DeepCompositionPoly::new(z, deep_coefficients);

        // combine all trace polynomials together and merge them into the DEEP composition
        // polynomial
        deep_composition_poly.add_trace_polys(trace_polys, ood_trace_states);

        // merge columns of constraint composition polynomial into the DEEP composition polynomial;
        deep_composition_poly.add_composition_poly(composition_poly, ood_evaluations);

        #[cfg(feature = "std")]
        debug!(
            "Built DEEP composition polynomial of degree {} in {} ms",
            deep_composition_poly.degree(),
            now.elapsed().as_millis()
        );

        // make sure the degree of the DEEP composition polynomial is equal to trace polynomial
        // degree minus 1.
        assert_eq!(domain.trace_length() - 2, deep_composition_poly.degree());

        // 5 ----- evaluate DEEP composition polynomial over LDE domain ---------------------------
        #[cfg(feature = "std")]
        let now = Instant::now();
        let deep_evaluations = deep_composition_poly.evaluate(&domain);
        // we check the following condition in debug mode only because infer_degree is an expensive
        // operation
        debug_assert_eq!(
            domain.trace_length() - 2,
            infer_degree(&deep_evaluations, domain.offset())
        );
        #[cfg(feature = "std")]
        debug!(
            "Evaluated DEEP composition polynomial over LDE domain (2^{} elements) in {} ms",
            domain.lde_domain_size().ilog2(),
            now.elapsed().as_millis()
        );

        // 6 ----- compute FRI layers for the composition polynomial ------------------------------
        #[cfg(feature = "std")]
        let now = Instant::now();
        let mut fri_prover = FriProver::new(air.options().to_fri_options());
        fri_prover.build_layers(&mut channel, deep_evaluations);
        #[cfg(feature = "std")]
        debug!(
            "Computed {} FRI layers from composition polynomial evaluations in {} ms",
            fri_prover.num_layers(),
            now.elapsed().as_millis()
        );

        // 7 ----- determine query positions ------------------------------------------------------
        #[cfg(feature = "std")]
        let now = Instant::now();

        // apply proof-of-work to the query seed
        channel.grind_query_seed();

        // generate pseudo-random query positions
        let query_positions = channel.get_query_positions();
        #[cfg(feature = "std")]
        debug!(
            "Determined {} query positions in {} ms",
            query_positions.len(),
            now.elapsed().as_millis()
        );

        // 8 ----- build proof object -------------------------------------------------------------
        #[cfg(feature = "std")]
        let now = Instant::now();

        // generate FRI proof
        let fri_proof = fri_prover.build_proof(&query_positions);

        // query the execution trace at the selected position; for each query, we need the
        // state of the trace at that position + Merkle authentication path
        let trace_queries = trace_commitment.query(&query_positions);

        // query the constraint commitment at the selected positions; for each query, we need just
        // a Merkle authentication path. this is because constraint evaluations for each step are
        // merged into a single value and Merkle authentication paths contain these values already
        let constraint_queries = constraint_commitment.query(&query_positions);

        // build the proof object
        let proof = channel.build_proof(trace_queries, constraint_queries, fri_proof);
        #[cfg(feature = "std")]
        debug!("Built proof object in {} ms", now.elapsed().as_millis());

        Ok(proof)
    }

    /// Computes a low-degree extension (LDE) of the provided execution trace over the specified
    /// domain and build a commitment to the extended trace.
    ///
    /// The extension is performed by interpolating each column of the execution trace into a
    /// polynomial of degree = trace_length - 1, and then evaluating the polynomial over the LDE
    /// domain.
    ///
    /// Trace commitment is computed by hashing each row of the extended execution trace, and then
    /// building a Merkle tree from the resulting hashes.
    fn build_trace_commitment<E>(
        &self,
        trace: &ColMatrix<E>,
        domain: &StarkDomain<Self::BaseField>,
    ) -> (RowMatrix<E>, MerkleTree<Self::HashFn>, ColMatrix<E>)
    where
        E: FieldElement<BaseField = Self::BaseField>,
    {
        // extend the execution trace
        #[cfg(feature = "std")]
        let now = Instant::now();
        let trace_polys = trace.interpolate_columns();
        let trace_lde =
            RowMatrix::evaluate_polys_over::<DEFAULT_SEGMENT_WIDTH>(&trace_polys, domain);
        #[cfg(feature = "std")]
        debug!(
            "Extended execution trace of {} columns from 2^{} to 2^{} steps ({}x blowup) in {} ms",
            trace_lde.num_cols(),
            trace_polys.num_rows().ilog2(),
            trace_lde.num_rows().ilog2(),
            domain.trace_to_lde_blowup(),
            now.elapsed().as_millis()
        );

        // build trace commitment
        #[cfg(feature = "std")]
        let now = Instant::now();
        let trace_tree = trace_lde.commit_to_rows();
        #[cfg(feature = "std")]
        debug!(
            "Computed execution trace commitment (Merkle tree of depth {}) in {} ms",
            trace_tree.depth(),
            now.elapsed().as_millis()
        );

        (trace_lde, trace_tree, trace_polys)
    }

    /// Evaluates constraint composition polynomial over the LDE domain and builds a commitment
    /// to these evaluations.
    ///
    /// The evaluation is done by evaluating each composition polynomial column over the LDE
    /// domain.
    ///
    /// The commitment is computed by hashing each row in the evaluation matrix, and then building
    /// a Merkle tree from the resulting hashes.
    fn build_constraint_commitment<E>(
        &self,
        composition_poly: &CompositionPoly<E>,
        domain: &StarkDomain<Self::BaseField>,
    ) -> ConstraintCommitment<E, Self::HashFn>
    where
        E: FieldElement<BaseField = Self::BaseField>,
    {
        // evaluate composition polynomial columns over the LDE domain
        #[cfg(feature = "std")]
        let now = Instant::now();
        let composed_evaluations = RowMatrix::evaluate_polys_over::<DEFAULT_SEGMENT_WIDTH>(
            composition_poly.data(),
            domain,
        );
        #[cfg(feature = "std")]
        debug!(
            "Evaluated {} composition polynomial columns over LDE domain (2^{} elements) in {} ms",
            composed_evaluations.num_cols(),
            composed_evaluations.num_rows().ilog2(),
            now.elapsed().as_millis()
        );

        // build constraint evaluation commitment
        #[cfg(feature = "std")]
        let now = Instant::now();
        let commitment = composed_evaluations.commit_to_rows();
        let constraint_commitment = ConstraintCommitment::new(composed_evaluations, commitment);
        #[cfg(feature = "std")]
        debug!(
            "Computed constraint evaluation commitment (Merkle tree of depth {}) in {} ms",
            constraint_commitment.tree_depth(),
            now.elapsed().as_millis()
        );
        constraint_commitment
    }
}


// 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::{constraints::CompositionPoly, StarkDomain, TracePolyTable};
use air::DeepCompositionCoefficients;
use math::{add_in_place, fft, mul_acc, polynom, ExtensionOf, FieldElement, StarkField};
use utils::{collections::Vec, iter_mut};

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

// DEEP COMPOSITION POLYNOMIAL
// ================================================================================================
pub struct DeepCompositionPoly<E: FieldElement> {
    coefficients: Vec<E>,
    cc: DeepCompositionCoefficients<E>,
    z: E,
}

impl<E: FieldElement> DeepCompositionPoly<E> {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new DEEP composition polynomial. Initially, this polynomial will be empty, and
    /// the intent is to populate the coefficients via add_trace_polys() and add_constraint_polys()
    /// methods.
    pub fn new(z: E, cc: DeepCompositionCoefficients<E>) -> Self {
        DeepCompositionPoly {
            coefficients: vec![],
            cc,
            z,
        }
    }

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

    /// Returns the size of the DEEP composition polynomial.
    pub fn poly_size(&self) -> usize {
        self.coefficients.len()
    }

    /// Returns the degree of the composition polynomial.
    pub fn degree(&self) -> usize {
        polynom::degree_of(&self.coefficients)
    }

    // TRACE POLYNOMIAL COMPOSITION
    // --------------------------------------------------------------------------------------------
    /// Combines all trace polynomials into a single polynomial and saves the result into
    /// the DEEP composition polynomial. The combination is done as follows:
    ///
    /// - Compute polynomials T'_i(x) = (T_i(x) - T_i(z)) / (x - z) and
    ///   T''_i(x) = (T_i(x) - T_i(z * g)) / (x - z * g) for all i, where T_i(x) is a trace
    ///   polynomial for column i.
    /// - Then, combine together all T'_i(x) and T''_i(x) polynomials using a random linear
    ///   combination as T(x) = sum((T'_i(x) + T''_i(x)) * cc_i) for all i, where cc_i is
    ///   the coefficient for the random linear combination drawn from the public coin.
    ///
    /// Note that evaluations of T_i(z) and T_i(z * g) are passed in via the `ood_trace_state`
    /// parameter.
    pub fn add_trace_polys(
        &mut self,
        trace_polys: TracePolyTable<E>,
        ood_trace_states: Vec<Vec<E>>,
    ) {
        assert!(self.coefficients.is_empty());

        // compute a second out-of-domain point offset from z by exactly trace generator; this
        // point defines the "next" computation state in relation to point z
        let trace_length = trace_polys.poly_size();
        let g = E::from(E::BaseField::get_root_of_unity(trace_length.ilog2()));
        let next_z = self.z * g;

        // combine trace polynomials into 2 composition polynomials T'(x) and T''(x)
        let mut t1_composition = E::zeroed_vector(trace_length);
        let mut t2_composition = E::zeroed_vector(trace_length);

        // index of a trace polynomial; we declare it here so that we can maintain index continuity
        // across all trace segments
        let mut i = 0;

        // --- merge polynomials of the main trace segment ----------------------------------------
        for poly in trace_polys.main_trace_polys() {
            // compute T'(x) = T(x) - T(z), multiply it by a pseudo-random coefficient,
            // and add the result into composition polynomial
            acc_trace_poly::<E::BaseField, E>(
                &mut t1_composition,
                poly,
                ood_trace_states[0][i],
                self.cc.trace[i],
            );

            // compute T''(x) = T(x) - T(z * g), multiply it by a pseudo-random coefficient,
            // and add the result into composition polynomial
            acc_trace_poly::<E::BaseField, E>(
                &mut t2_composition,
                poly,
                ood_trace_states[1][i],
                self.cc.trace[i],
            );

            i += 1;
        }

        // --- merge polynomials of the auxiliary trace segments ----------------------------------
        for poly in trace_polys.aux_trace_polys() {
            // compute T'(x) = T(x) - T(z), multiply it by a pseudo-random coefficient,
            // and add the result into composition polynomial
            acc_trace_poly::<E, E>(
                &mut t1_composition,
                poly,
                ood_trace_states[0][i],
                self.cc.trace[i],
            );

            // compute T''(x) = T(x) - T(z * g), multiply it by a pseudo-random coefficient,
            // and add the result into composition polynomial
            acc_trace_poly::<E, E>(
                &mut t2_composition,
                poly,
                ood_trace_states[1][i],
                self.cc.trace[i],
            );

            i += 1;
        }

        // divide the composition polynomials by (x - z) and (x - z * g), respectively,
        // and add the resulting polynomials together; the output of this step
        // is a single trace polynomial T(x) and deg(T(x)) = trace_length - 2.
        let trace_poly =
            merge_trace_compositions(vec![t1_composition, t2_composition], vec![self.z, next_z]);

        // set the coefficients of the DEEP composition polynomial
        self.coefficients = trace_poly;
        assert_eq!(self.poly_size() - 2, self.degree());
    }

    // CONSTRAINT POLYNOMIAL COMPOSITION
    // --------------------------------------------------------------------------------------------
    /// Divides out OOD point z from the constraint composition polynomial and saves the result
    /// into the DEEP composition polynomial. This method is intended to be called only after the
    /// add_trace_polys() method has been executed. The composition is done as follows:
    ///
    /// - For each H_i(x), compute H'_i(x) = (H_i(x) - H(z)) / (x - z), where H_i(x) is the
    ///   ith composition polynomial column.
    /// - Then, combine all H_i(x) polynomials together by computing H(x) = sum(H_i(x) * cc_i) for
    ///   all i, where cc_i is the coefficient for the random linear combination drawn from the
    ///   public coin.
    ///
    /// Note that evaluations of H_i(x) at z are passed in via the `ood_evaluations` parameter.
    pub fn add_composition_poly(
        &mut self,
        composition_poly: CompositionPoly<E>,
        ood_evaluations: Vec<E>,
    ) {
        assert!(!self.coefficients.is_empty());

        let z = self.z;

        let mut column_polys = composition_poly.into_columns();

        // Divide out the OOD point z from column polynomials
        iter_mut!(column_polys)
            .zip(ood_evaluations)
            .for_each(|(poly, value_at_z)| {
                // compute H'_i(x) = (H_i(x) - H_i(z)) / (x - z)
                poly[0] -= value_at_z;
                polynom::syn_div_in_place(poly, 1, z);
            });

        // add H'_i(x) * cc_i for all i into the DEEP composition polynomial
        for (i, poly) in column_polys.into_iter().enumerate() {
            mul_acc::<E, E>(&mut self.coefficients, &poly, self.cc.constraints[i]);
        }
        assert_eq!(self.poly_size() - 2, self.degree());
    }

    // LOW-DEGREE EXTENSION
    // --------------------------------------------------------------------------------------------
    /// Evaluates DEEP composition polynomial over the specified LDE domain and returns the result.
    pub fn evaluate(self, domain: &StarkDomain<E::BaseField>) -> Vec<E> {
        fft::evaluate_poly_with_offset(
            &self.coefficients,
            domain.trace_twiddles(),
            domain.offset(),
            domain.trace_to_lde_blowup(),
        )
    }
}

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

/// Divides each polynomial in the list by the corresponding divisor, and computes the
/// coefficient-wise sum of all resulting polynomials.
fn merge_trace_compositions<E: FieldElement>(mut polys: Vec<Vec<E>>, divisors: Vec<E>) -> Vec<E> {
    // divide all polynomials by their corresponding divisor
    iter_mut!(polys).zip(divisors).for_each(|(poly, divisor)| {
        polynom::syn_div_in_place(poly, 1, divisor);
    });

    // add all polynomials together into a single polynomial
    let mut result = polys.remove(0);
    for poly in polys.iter() {
        add_in_place(&mut result, poly);
    }

    result
}

/// Computes (P(x) - value) * k and saves the result into the accumulator.
fn acc_trace_poly<F, E>(accumulator: &mut [E], poly: &[F], value: E, k: E)
where
    F: FieldElement,
    E: FieldElement<BaseField = F::BaseField> + ExtensionOf<F>,
{
    mul_acc(accumulator, poly, k);
    let adjusted_tz = value * k;
    accumulator[0] -= adjusted_tz;
}