// 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::fft_inputs::FftInputs;
use crate::field::{FieldElement, StarkField};
use utils::{collections::Vec, iterators::*, rayon, uninit_vector};

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

/// Evaluates polynomial `p` using FFT algorithm; the evaluation is done in-place, meaning
/// `p` is updated with results of the evaluation.
pub fn evaluate_poly<B: StarkField, E: FieldElement<BaseField = B>>(p: &mut [E], twiddles: &[B]) {
    split_radix_fft(p, twiddles);
    permute(p);
}

/// Evaluates polynomial `p` using FFT algorithm and returns the result. The polynomial is
/// evaluated over domain specified by `twiddles`, expanded by the `blowup_factor`, and shifted
/// by the `domain_offset`.
pub fn evaluate_poly_with_offset<B: StarkField, E: FieldElement<BaseField = B>>(
    p: &[E],
    twiddles: &[B],
    domain_offset: B,
    blowup_factor: usize,
) -> Vec<E> {
    let domain_size = p.len() * blowup_factor;
    let g = B::get_root_of_unity(domain_size.ilog2());
    let mut result = unsafe { uninit_vector(domain_size) };

    result
        .as_mut_slice()
        .par_chunks_mut(p.len())
        .enumerate()
        .for_each(|(i, chunk)| {
            let idx = super::permute_index(blowup_factor, i) as u64;
            let offset = g.exp(idx.into()) * domain_offset;
            clone_and_shift(p, chunk, offset);
            split_radix_fft(chunk, twiddles);
        });

    permute(&mut result);
    result
}

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

/// Uses FFT algorithm to interpolate a polynomial from provided `values`; the interpolation
/// is done in-place, meaning `values` are updated with polynomial coefficients.
pub fn interpolate_poly<B, E>(v: &mut [E], inv_twiddles: &[B])
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    split_radix_fft(v, inv_twiddles);
    let inv_length = E::inv((v.len() as u64).into());
    v.par_iter_mut().for_each(|e| *e *= inv_length);
    permute(v);
}

/// Uses FFT algorithm to interpolate a polynomial from provided `values` over the domain defined
/// by `inv_twiddles` and offset by `domain_offset` factor.
pub fn interpolate_poly_with_offset<B, E>(values: &mut [E], inv_twiddles: &[B], domain_offset: B)
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    split_radix_fft(values, inv_twiddles);
    permute(values);

    let domain_offset = E::inv(domain_offset.into());
    let inv_len = E::inv((values.len() as u64).into());
    let batch_size = values.len() / rayon::current_num_threads().next_power_of_two();

    values
        .par_chunks_mut(batch_size)
        .enumerate()
        .for_each(|(i, batch)| {
            let mut offset = domain_offset.exp(((i * batch_size) as u64).into()) * inv_len;
            for coeff in batch.iter_mut() {
                *coeff = *coeff * offset;
                offset = offset * domain_offset;
            }
        });
}

// PERMUTATIONS
// ================================================================================================

pub fn permute<E: FieldElement>(v: &mut [E]) {
    let n = v.len();
    let num_batches = rayon::current_num_threads().next_power_of_two();
    let batch_size = n / num_batches;
    rayon::scope(|s| {
        for batch_idx in 0..num_batches {
            // create another mutable reference to the slice of values to use in a new thread; this
            // is OK because we never write the same positions in the slice from different threads
            let values = unsafe { &mut *(&mut v[..] as *mut [E]) };
            s.spawn(move |_| {
                let batch_start = batch_idx * batch_size;
                let batch_end = batch_start + batch_size;
                for i in batch_start..batch_end {
                    let j = super::permute_index(n, i);
                    if j > i {
                        values.swap(i, j);
                    }
                }
            });
        }
    });
}

// SPLIT-RADIX FFT
// ================================================================================================

/// In-place recursive FFT with permuted output.
/// Adapted from: https://github.com/0xProject/OpenZKP/tree/master/algebra/primefield/src/fft
pub(super) fn split_radix_fft<B: StarkField, E: FieldElement<BaseField = B>>(
    values: &mut [E],
    twiddles: &[B],
) {
    // generator of the domain should be in the middle of twiddles
    let n = values.len();
    let g = twiddles[twiddles.len() / 2];
    debug_assert_eq!(g.exp((n as u32).into()), E::BaseField::ONE);

    let inner_len = 1_usize << (n.ilog2() / 2);
    let outer_len = n / inner_len;
    let stretch = outer_len / inner_len;
    debug_assert!(outer_len == inner_len || outer_len == 2 * inner_len);
    debug_assert_eq!(outer_len * inner_len, n);

    // transpose inner x inner x stretch square matrix
    transpose_square_stretch(values, inner_len, stretch);

    // apply inner FFTs
    values
        .par_chunks_mut(outer_len)
        .for_each(|row| row.fft_in_place_raw(&twiddles, stretch, stretch, 0));

    // transpose inner x inner x stretch square matrix
    transpose_square_stretch(values, inner_len, stretch);

    // apply outer FFTs
    values
        .par_chunks_mut(outer_len)
        .enumerate()
        .for_each(|(i, row)| {
            if i > 0 {
                let i = super::permute_index(inner_len, i);
                let inner_twiddle = g.exp((i as u32).into());
                let mut outer_twiddle = inner_twiddle;
                for element in row.iter_mut().skip(1) {
                    *element = (*element).mul_base(outer_twiddle);
                    outer_twiddle = outer_twiddle * inner_twiddle;
                }
            }
            row.fft_in_place(&twiddles);
        });
}

// TRANSPOSING
// ================================================================================================

fn transpose_square_stretch<T>(matrix: &mut [T], size: usize, stretch: usize) {
    assert_eq!(matrix.len(), size * size * stretch);
    match stretch {
        1 => transpose_square_1(matrix, size),
        2 => transpose_square_2(matrix, size),
        _ => unimplemented!("only stretch sizes 1 and 2 are supported"),
    }
}

fn transpose_square_1<T>(matrix: &mut [T], size: usize) {
    debug_assert_eq!(matrix.len(), size * size);
    if size % 2 != 0 {
        unimplemented!("odd sizes are not supported");
    }

    // iterate over upper-left triangle, working in 2x2 blocks
    for row in (0..size).step_by(2) {
        let i = row * size + row;
        matrix.swap(i + 1, i + size);
        for col in (row..size).step_by(2).skip(1) {
            let i = row * size + col;
            let j = col * size + row;
            matrix.swap(i, j);
            matrix.swap(i + 1, j + size);
            matrix.swap(i + size, j + 1);
            matrix.swap(i + size + 1, j + size + 1);
        }
    }
}

fn transpose_square_2<T>(matrix: &mut [T], size: usize) {
    debug_assert_eq!(matrix.len(), 2 * size * size);

    // iterate over upper-left triangle, working in 1x2 blocks
    for row in 0..size {
        for col in (row..size).skip(1) {
            let i = (row * size + col) * 2;
            let j = (col * size + row) * 2;
            matrix.swap(i, j);
            matrix.swap(i + 1, j + 1);
        }
    }
}

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

fn clone_and_shift<E: FieldElement>(source: &[E], destination: &mut [E], offset: E::BaseField) {
    let batch_size = source.len() / rayon::current_num_threads().next_power_of_two();
    source
        .par_chunks(batch_size)
        .zip(destination.par_chunks_mut(batch_size))
        .enumerate()
        .for_each(|(i, (source, destination))| {
            let mut factor = offset.exp(((i * batch_size) as u64).into());
            for (s, d) in source.iter().zip(destination.iter_mut()) {
                *d = (*s).mul_base(factor);
                factor = factor * 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.

use super::{permute_index, FieldElement};

// CONSTANTS
// ================================================================================================
const MAX_LOOP: usize = 256;

// FFT INPUTS TRAIT
// ================================================================================================

/// Defines the interface that must be implemented by the input to fft_in_place method.
#[allow(clippy::len_without_is_empty)]
pub trait FftInputs<E: FieldElement> {
    /// Returns the number of elements in this input.
    fn len(&self) -> usize;

    /// Combines the result of smaller number theoretic transform into a larger NTT.
    fn butterfly(&mut self, offset: usize, stride: usize);

    /// Combines the result of smaller number theoretic transform multiplied with a
    /// twiddle factor into a larger NTT.
    fn butterfly_twiddle(&mut self, twiddle: E::BaseField, offset: usize, stride: usize);

    /// Swaps the element at index i with the element at index j. Specifically:
    ///
    /// elem_i <-> elem_j
    ///
    /// # Panics
    /// Panics if i or j are out of bounds.
    fn swap(&mut self, i: usize, j: usize);

    /// Multiplies every element in this input by a series of increment. Specifically:
    ///
    /// elem_i = elem_i * offset * increment^i
    fn shift_by_series(&mut self, offset: E::BaseField, increment: E::BaseField);

    /// Multiplies every element in this input by `offset`. Specifically:
    ///
    /// elem_i = elem_i * offset
    fn shift_by(&mut self, offset: E::BaseField);

    /// Permutes the elements in this input using the permutation defined by the given
    /// permutation index.
    ///
    /// The permutation index is a number between 0 and `self.len() - 1` that specifies the
    /// permutation to apply to the input. The permutation is applied in place, so the input
    /// is replaced with the result of the permutation. The permutation is applied by swapping
    /// elements in the input.
    ///
    /// # Panics
    /// Panics if the permutation index is out of bounds.
    fn permute(&mut self) {
        let n = self.len();
        for i in 0..n {
            let j = permute_index(n, i);
            if j > i {
                self.swap(i, j);
            }
        }
    }

    /// Applies the FFT to this input.
    ///
    /// The FFT is applied in place, so the input is replaced with the result of the FFT. The
    /// `twiddles` parameter specifies the twiddle factors to use for the FFT.
    ///
    /// This is a convenience method equivalent to calling fft_in_place_raw(twiddles, 1, 1, 0).
    ///
    /// # Panics
    /// Panics if length of the `twiddles` parameter is not self.len() / 2.
    fn fft_in_place(&mut self, twiddles: &[E::BaseField]) {
        fft_in_place(self, twiddles, 1, 1, 0);
    }

    /// Applies the FFT to this input.
    ///
    /// The FFT is applied in place, so the input is replaced with the result of the FFT. The
    /// `twiddles` parameter specifies the twiddle factors to use for the FFT.
    ///
    /// # Panics
    /// Panics if length of the `twiddles` parameter is not self.len() / 2.
    fn fft_in_place_raw(
        &mut self,
        twiddles: &[E::BaseField],
        count: usize,
        stride: usize,
        offset: usize,
    ) {
        fft_in_place(self, twiddles, count, stride, offset)
    }
}

// SLICE IMPLEMENTATION
// ================================================================================================

/// Implements FftInputs for a slice of field elements.
impl<E: FieldElement> FftInputs<E> for [E] {
    fn len(&self) -> usize {
        self.len()
    }

    #[inline(always)]
    fn butterfly(&mut self, offset: usize, stride: usize) {
        let i = offset;
        let j = offset + stride;
        let temp = self[i];
        self[i] = temp + self[j];
        self[j] = temp - self[j];
    }

    #[inline(always)]
    fn butterfly_twiddle(&mut self, twiddle: E::BaseField, offset: usize, stride: usize) {
        let i = offset;
        let j = offset + stride;
        let temp = self[i];
        self[j] = self[j].mul_base(twiddle);
        self[i] = temp + self[j];
        self[j] = temp - self[j];
    }

    fn swap(&mut self, i: usize, j: usize) {
        self.swap(i, j)
    }

    fn shift_by_series(&mut self, offset: E::BaseField, increment: E::BaseField) {
        let mut offset = E::from(offset);
        let increment = E::from(increment);
        for d in self.iter_mut() {
            *d *= offset;
            offset *= increment;
        }
    }

    fn shift_by(&mut self, offset: E::BaseField) {
        let offset = E::from(offset);
        for d in self.iter_mut() {
            *d *= offset;
        }
    }
}

// SLICE OF ARRAYS IMPLEMENTATION
// ================================================================================================

/// Implements [FftInputs] for a slice of field element arrays.
#[allow(clippy::needless_range_loop)]
impl<E: FieldElement, const N: usize> FftInputs<E> for [[E; N]] {
    fn len(&self) -> usize {
        self.len()
    }

    #[inline(always)]
    fn butterfly(&mut self, offset: usize, stride: usize) {
        let i = offset;
        let j = offset + stride;

        let temp = self[i];
        for col_idx in 0..N {
            self[i][col_idx] = temp[col_idx] + self[j][col_idx];
            self[j][col_idx] = temp[col_idx] - self[j][col_idx];
        }
    }

    #[inline(always)]
    fn butterfly_twiddle(&mut self, twiddle: E::BaseField, offset: usize, stride: usize) {
        let i = offset;
        let j = offset + stride;

        let twiddle = E::from(twiddle);
        let temp = self[i];

        for col_idx in 0..N {
            self[j][col_idx] *= twiddle;
            self[i][col_idx] = temp[col_idx] + self[j][col_idx];
            self[j][col_idx] = temp[col_idx] - self[j][col_idx];
        }
    }

    fn swap(&mut self, i: usize, j: usize) {
        self.swap(i, j)
    }

    fn shift_by(&mut self, offset: E::BaseField) {
        let offset = E::from(offset);
        for row_idx in 0..self.len() {
            for col_idx in 0..N {
                self[row_idx][col_idx] *= offset;
            }
        }
    }

    fn shift_by_series(&mut self, offset: E::BaseField, increment: E::BaseField) {
        let increment = E::from(increment);
        let mut offset = E::from(offset);

        for row_idx in 0..self.len() {
            for col_idx in 0..N {
                self[row_idx][col_idx] *= offset;
            }
            offset *= increment;
        }
    }
}

// CORE FFT ALGORITHM
// ================================================================================================

/// In-place recursive FFT with permuted output.
///
/// Adapted from: https://github.com/0xProject/OpenZKP/tree/master/algebra/primefield/src/fft
fn fft_in_place<E, I>(
    values: &mut I,
    twiddles: &[E::BaseField],
    count: usize,
    stride: usize,
    offset: usize,
) where
    E: FieldElement,
    I: FftInputs<E> + ?Sized,
{
    let size = values.len() / stride;
    debug_assert!(size.is_power_of_two());
    debug_assert!(offset < stride);
    debug_assert_eq!(values.len() % size, 0);

    // Keep recursing until size is 2
    if size > 2 {
        if stride == count && count < MAX_LOOP {
            fft_in_place(values, twiddles, 2 * count, 2 * stride, offset);
        } else {
            fft_in_place(values, twiddles, count, 2 * stride, offset);
            fft_in_place(values, twiddles, count, 2 * stride, offset + stride);
        }
    }

    // Apply butterfly operations.
    for offset in offset..(offset + count) {
        I::butterfly(values, offset, stride);
    }

    // Apply butterfly operations with twiddle factors.
    let last_offset = offset + size * stride;
    for (i, offset) in (offset..last_offset)
        .step_by(2 * stride)
        .enumerate()
        .skip(1)
    {
        for j in offset..(offset + count) {
            I::butterfly_twiddle(values, twiddles[i], j, stride);
        }
    }
}


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

//! FFT-based polynomial evaluation and interpolation.
//!
//! Functions in this module can be used to evaluate and interpolate polynomials over domains
//! which are multiplicative subgroups of finite fields and have lengths equal to powers of two.
//! As compared to evaluation and interpolation functions available in the `polynom` module,
//! these functions are much more efficient: their runtime complexity is O(`n` log `n`), where
//! `n` is the domain size.

use crate::{
    fft::fft_inputs::FftInputs,
    field::{FieldElement, StarkField},
    utils::get_power_series,
};

pub mod fft_inputs;
pub mod real_u64;
mod serial;

#[cfg(feature = "concurrent")]
mod concurrent;

use utils::collections::Vec;

#[cfg(test)]
mod tests;

// CONSTANTS
// ================================================================================================
const MIN_CONCURRENT_SIZE: usize = 1024;

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

/// Evaluates a polynomial on all points of the specified domain using the FFT algorithm.
///
/// Uses the [FFT](https://en.wikipedia.org/wiki/Discrete_Fourier_transform_(general)) algorithm
/// to evaluate polynomial `p` on all points of a domain defined by the length of `p` in the field
/// specified by the `B` type parameter. The evaluation is done in-place, meaning no additional
/// memory is allocated and `p` is updated with results of the evaluation. The polynomial `p`
/// is expected to be in coefficient form.
///
/// The complexity of evaluation is O(`n` log(`n`)), where `n` is the size of the domain.
///
/// The size of the domain is assumed to be equal to `p.len()` which must be a power of two. The
/// base field specified by `B` must have a multiplicative subgroup of size equal to `p.len()`.
///
/// The `twiddles` needed for evaluation can be obtained via `fft::get_twiddles()` function using
/// `p.len()` as the domain size parameter. This implies that `twiddles.len()` must be equal to
/// `p.len()` / 2.
///
/// When `concurrent` feature is enabled, the evaluation is done in multiple threads.
///
/// # Panics
/// Panics if:
/// * Length of `p` is not a power of two.
/// * Length of `twiddles` is not `p.len()` / 2.
/// * Field specified by `B` does not contain a multiplicative subgroup of size `p.len()`.
///
/// # Examples
/// ```
/// # use winter_math::{polynom, fft::*, get_power_series};
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement, StarkField};
/// # use rand_utils::rand_vector;
/// let n = 2048;
///
/// // build a random polynomial
/// let mut p: Vec<BaseElement> = rand_vector(n);
///
/// // evaluate the polynomial over the domain using regular polynomial evaluation
/// let g = BaseElement::get_root_of_unity(n.ilog2());
/// let domain = get_power_series(g, n);
/// let expected = polynom::eval_many(&p, &domain);
///
/// // evaluate the polynomial over the domain using FFT-based evaluation
/// let twiddles = get_twiddles::<BaseElement>(p.len());
/// evaluate_poly(&mut p, &twiddles);
///
/// assert_eq!(expected, p);
/// ```
pub fn evaluate_poly<B, E>(p: &mut [E], twiddles: &[B])
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    assert!(
        p.len().is_power_of_two(),
        "number of coefficients must be a power of 2"
    );
    assert_eq!(
        p.len(),
        twiddles.len() * 2,
        "invalid number of twiddles: expected {} but received {}",
        p.len() / 2,
        twiddles.len()
    );
    assert!(
        p.len().ilog2() <= B::TWO_ADICITY,
        "multiplicative subgroup of size {} does not exist in the specified base field",
        p.len()
    );

    // when `concurrent` feature is enabled, run the concurrent version of the function; unless
    // the polynomial is small, then don't bother with the concurrent version
    if cfg!(feature = "concurrent") && p.len() >= MIN_CONCURRENT_SIZE {
        #[cfg(feature = "concurrent")]
        concurrent::evaluate_poly(p, twiddles);
    } else {
        serial::evaluate_poly(p, twiddles);
    }
}

/// Evaluates a polynomial on all points of the specified (shifted) domain using the FFT algorithm.
///
/// Uses the [FFT](https://en.wikipedia.org/wiki/Discrete_Fourier_transform_(general)) algorithm
/// to evaluate polynomial `p` on all points of a domain defined by the length of `p`, expanded
/// by the `blowup_factor`, and shifted by the `domain_offset` in the field specified by the `B`
/// type parameter. The polynomial `p` is expected to be in coefficient form.
///
/// The complexity of evaluation is O(`n` log(`n`)), where `n` is the size of the domain.
///
/// The size of the domain is assumed to be equal to `p.len()` * `blowup_factor` both of which must
/// be powers of two. The base field specified by `B` must have a multiplicative subgroup of size
/// equal to `p.len()` * `blowup_factor`.
///
/// The shifted domain is defined as the original domain with every element multiplied by the
/// `domain_offset`.
///
/// The `twiddles` needed for evaluation can be obtained via `fft::get_twiddles()` function using
/// `p.len()` as the domain size parameter. This implies that `twiddles.len()` must be equal to
/// `p.len()` / 2.
///
/// When `concurrent` feature is enabled, the evaluation is done in multiple threads.
///
/// # Panics
/// Panics if:
/// * Length of `p` is not a power of two.
/// * `blowup_factor` is not a power of two.
/// * Length of `twiddles` is not `p.len()` / 2.
/// * Field specified by `B` does not contain a multiplicative subgroup of size `p.len()`.
/// * `domain_offset` is ZERO.
///
/// # Examples
/// ```
/// # use winter_math::{polynom, fft::*, get_power_series};
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement, StarkField};
/// # use rand_utils::rand_vector;
/// let n = 2048;
/// let offset = BaseElement::GENERATOR;
/// let blowup_factor = 2;
///
/// // build a random polynomial
/// let mut p: Vec<BaseElement> = rand_vector(n / blowup_factor);
///
/// // evaluate the polynomial over the domain using regular polynomial evaluation
/// let g = BaseElement::get_root_of_unity(n.ilog2());
/// let domain = get_power_series(g, n);
/// let shifted_domain = domain.iter().map(|&x| x * offset).collect::<Vec<_>>();
/// let expected = polynom::eval_many(&p, &shifted_domain);
///
/// // evaluate the polynomial over the domain using FFT-based evaluation
/// let twiddles = get_twiddles::<BaseElement>(p.len());
/// let actual = evaluate_poly_with_offset(&mut p, &twiddles, offset, blowup_factor);
///
/// assert_eq!(expected, actual);
/// ```
pub fn evaluate_poly_with_offset<B, E>(
    p: &[E],
    twiddles: &[B],
    domain_offset: B,
    blowup_factor: usize,
) -> Vec<E>
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    assert!(
        p.len().is_power_of_two(),
        "number of coefficients must be a power of 2"
    );
    assert!(
        blowup_factor.is_power_of_two(),
        "blowup factor must be a power of 2"
    );
    assert_eq!(
        p.len(),
        twiddles.len() * 2,
        "invalid number of twiddles: expected {} but received {}",
        p.len() / 2,
        twiddles.len()
    );
    assert!(
        (p.len() * blowup_factor).ilog2() <= B::TWO_ADICITY,
        "multiplicative subgroup of size {} does not exist in the specified base field",
        p.len() * blowup_factor
    );
    assert_ne!(domain_offset, B::ZERO, "domain offset cannot be zero");

    // assign a dummy value here to make the compiler happy
    #[allow(unused_assignments)]
    let mut result = Vec::new();

    // when `concurrent` feature is enabled, run the concurrent version of the function; unless
    // the polynomial is small, then don't bother with the concurrent version
    if cfg!(feature = "concurrent") && p.len() >= MIN_CONCURRENT_SIZE {
        #[cfg(feature = "concurrent")]
        {
            result =
                concurrent::evaluate_poly_with_offset(p, twiddles, domain_offset, blowup_factor);
        }
    } else {
        result = serial::evaluate_poly_with_offset(p, twiddles, domain_offset, blowup_factor);
    }

    result
}

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

/// Interpolates evaluations of a polynomial over the specified domain into a polynomial in
/// coefficient from using the FFT algorithm.
///
/// Uses the inverse [FFT](https://en.wikipedia.org/wiki/Discrete_Fourier_transform_(general))
/// algorithm to interpolate a polynomial from its evaluations over a domain defined by the
/// length of `evaluations` in the field specified by the `B` type parameter.  The interpolation
/// is done in-place, meaning no additional memory is allocated and the evaluations contained in
/// `evaluations` are replaced with polynomial coefficients.
///
/// The complexity of interpolation is O(`n` log(`n`)), where `n` is the size of the domain.
///
/// The size of the domain is assumed to be equal to `evaluations.len()` which must be a power
/// of two. The base field specified by `B` must have a multiplicative subgroup of size equal
/// to `evaluations.len()`.
///
/// The `inv_twiddles` needed for interpolation can be obtained via `fft::get_inv_twiddles()`
/// function using `evaluations.len()` as the domain size parameter. This implies that
/// `twiddles.len()` must be equal to `evaluations.len()` / 2.
///
/// When `concurrent` feature is enabled, the interpolation is done in multiple threads.
///
/// # Panics
/// Panics if:
/// * Length of `evaluations` is not a power of two.
/// * Length of `inv_twiddles` is not `evaluations.len()` / 2.
/// * Field specified by `B` does not contain a multiplicative subgroup of size
///   `evaluations.len()`.
///
/// # Examples
/// ```
/// # use winter_math::{polynom, fft::*, get_power_series};
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement, StarkField};
/// # use rand_utils::rand_vector;
/// let n = 2048;
///
/// // build a random polynomial
/// let p: Vec<BaseElement> = rand_vector(n);
///
/// // evaluate the polynomial over the domain using regular polynomial evaluation
/// let g = BaseElement::get_root_of_unity(n.ilog2());
/// let domain = get_power_series(g, n);
/// let mut ys = polynom::eval_many(&p, &domain);
///
/// // interpolate the evaluations into a polynomial
/// let inv_twiddles = get_inv_twiddles::<BaseElement>(ys.len());
/// interpolate_poly(&mut ys, &inv_twiddles);
///
/// assert_eq!(p, ys);
/// ```
pub fn interpolate_poly<B, E>(evaluations: &mut [E], inv_twiddles: &[B])
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    assert!(
        evaluations.len().is_power_of_two(),
        "number of evaluations must be a power of 2, but was {}",
        evaluations.len()
    );
    assert_eq!(
        evaluations.len(),
        inv_twiddles.len() * 2,
        "invalid number of twiddles: expected {} but received {}",
        evaluations.len() / 2,
        inv_twiddles.len()
    );
    assert!(
        evaluations.len().ilog2() <= B::TWO_ADICITY,
        "multiplicative subgroup of size {} does not exist in the specified base field",
        evaluations.len()
    );

    // when `concurrent` feature is enabled, run the concurrent version of interpolate_poly;
    // unless the number of evaluations is small, then don't bother with the concurrent version
    if cfg!(feature = "concurrent") && evaluations.len() >= MIN_CONCURRENT_SIZE {
        #[cfg(feature = "concurrent")]
        concurrent::interpolate_poly(evaluations, inv_twiddles);
    } else {
        serial::interpolate_poly(evaluations, inv_twiddles);
    }
}

/// Interpolates evaluations of a polynomial over the specified (shifted) domain into a polynomial
/// in coefficient from using the FFT algorithm.
///
/// Uses the inverse [FFT](https://en.wikipedia.org/wiki/Discrete_Fourier_transform_(general))
/// algorithm to interpolate a polynomial from its evaluations over a domain defined by the
/// length of `evaluations` and shifted by the `domain_offset` in the field specified by the `B`
/// type parameter. The interpolation is done in-place, meaning no additional memory is allocated
/// and the evaluations contained in `evaluations` are replaced with polynomial coefficients.
///
/// The complexity of interpolation is O(`n` log(`n`)), where `n` is the size of the domain.
///
/// The size of the domain is assumed to be equal to `evaluations.len()` which must be a power
/// of two. The base field specified by `B` must have a multiplicative subgroup of size equal
/// to `evaluations.len()`.
///
/// The shifted domain is defined as the original domain with every element multiplied by the
/// `domain_offset`.
///
/// The `inv_twiddles` needed for interpolation can be obtained via `fft::get_inv_twiddles()`
/// function using `evaluations.len()` as the domain size parameter. This implies that
/// `twiddles.len()` must be equal to `evaluations.len()` / 2.
///
/// When `concurrent` feature is enabled, the interpolation is done in multiple threads.
///
/// # Panics
/// Panics if:
/// * Length of `evaluations` is not a power of two.
/// * Length of `inv_twiddles` is not `evaluations.len()` / 2.
/// * Field specified by `B` does not contain a multiplicative subgroup of size
///   `evaluations.len()`.
/// * `domain_offset` is ZERO.
///
/// # Examples
/// ```
/// # use winter_math::{polynom, fft::*, get_power_series};
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement, StarkField};
/// # use rand_utils::rand_vector;
/// let n = 2048;
/// let offset = BaseElement::GENERATOR;
///
/// // build a random polynomial
/// let p: Vec<BaseElement> = rand_vector(n);
///
/// // evaluate the polynomial over the domain using regular polynomial evaluation
/// let g = BaseElement::get_root_of_unity(n.ilog2());
/// let domain = get_power_series(g, n);
/// let shifted_domain = domain.iter().map(|&x| x * offset).collect::<Vec<_>>();
/// let mut ys = polynom::eval_many(&p, &shifted_domain);
///
/// // interpolate the evaluations into a polynomial
/// let inv_twiddles = get_inv_twiddles::<BaseElement>(ys.len());
/// interpolate_poly_with_offset(&mut ys, &inv_twiddles, offset);
///
/// assert_eq!(p, ys);
/// ```
pub fn interpolate_poly_with_offset<B, E>(
    evaluations: &mut [E],
    inv_twiddles: &[B],
    domain_offset: B,
) where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    assert!(
        evaluations.len().is_power_of_two(),
        "number of evaluations must be a power of 2, but was {}",
        evaluations.len()
    );
    assert_eq!(
        evaluations.len(),
        inv_twiddles.len() * 2,
        "invalid number of twiddles: expected {} but received {}",
        evaluations.len() / 2,
        inv_twiddles.len()
    );
    assert!(
        evaluations.len().ilog2() <= B::TWO_ADICITY,
        "multiplicative subgroup of size {} does not exist in the specified base field",
        evaluations.len()
    );
    assert_ne!(domain_offset, B::ZERO, "domain offset cannot be zero");

    // when `concurrent` feature is enabled, run the concurrent version of the function; unless
    // the polynomial is small, then don't bother with the concurrent version
    if cfg!(feature = "concurrent") && evaluations.len() >= MIN_CONCURRENT_SIZE {
        #[cfg(feature = "concurrent")]
        concurrent::interpolate_poly_with_offset(evaluations, inv_twiddles, domain_offset);
    } else {
        serial::interpolate_poly_with_offset(evaluations, inv_twiddles, domain_offset);
    }
}

// RAW FFT ALGORITHM
// ================================================================================================

/// Executes a single-threaded version of the FFT algorithm on the provided values.
///
/// The evaluation is done in-place, meaning the function does not allocate any additional memory,
/// and the results are written back into `values`.
///
/// The `twiddles` needed for evaluation can be obtained via `fft::get_twiddles()` function using
/// `values.len()` as the domain size parameter. This implies that `twiddles.len()` must be equal
/// to `values.len()` / 2.
///
/// # Panics
/// Panics if:
/// * Length of `values` is not a power of two.
/// * Length of `twiddles` is not `values.len()` / 2.
/// * Field specified by `B` does not contain a multiplicative subgroup of size `values.len()`.
pub fn serial_fft<B, E>(values: &mut [E], twiddles: &[B])
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    assert!(
        values.len().is_power_of_two(),
        "number of values must be a power of 2, but was {}",
        values.len()
    );
    assert_eq!(
        values.len(),
        twiddles.len() * 2,
        "invalid number of twiddles: expected {} but received {}",
        values.len() / 2,
        twiddles.len()
    );
    assert!(
        values.len().ilog2() <= B::TWO_ADICITY,
        "multiplicative subgroup of size {} does not exist in the specified base field",
        values.len()
    );
    values.fft_in_place(twiddles);
    values.permute();
}

// TWIDDLES
// ================================================================================================

/// Returns a set of twiddles for the specified domain size.
///
/// These twiddles can then be used for FFT-based polynomial evaluation. The length of the returned
/// vector will be equal to `domain_size` / 2.
///
/// When `concurrent` feature is enabled, the twiddles are generated in multiple threads.
///
/// # Panics
/// Panics if:
/// * `domain_size` is not a power of two.
/// * Field specified by `B` does not contain a multiplicative subgroup of size `domain_size`.
///
/// # Examples
/// ```
/// # use winter_math::fft::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement, StarkField};
/// let n = 2048;
/// let twiddles = get_twiddles::<BaseElement>(n);
///
/// assert_eq!(n / 2, twiddles.len());
/// ```
pub fn get_twiddles<B>(domain_size: usize) -> Vec<B>
where
    B: StarkField,
{
    assert!(
        domain_size.is_power_of_two(),
        "domain size must be a power of 2"
    );
    assert!(
        domain_size.ilog2() <= B::TWO_ADICITY,
        "multiplicative subgroup of size {domain_size} does not exist in the specified base field"
    );
    let root = B::get_root_of_unity(domain_size.ilog2());
    let mut twiddles = get_power_series(root, domain_size / 2);
    permute(&mut twiddles);
    twiddles
}

/// Returns a set of inverse twiddles for the specified domain size.
///
/// These twiddles can then be used for FFT-based polynomial interpolation. The length of the
/// returned vector will be equal to `domain_size` / 2.
///
/// When `concurrent` feature is enabled, the twiddles are generated in multiple threads.
///
/// # Panics
/// Panics if:
/// * `domain_size` is not a power of two.
/// * Field specified by `B` does not contain a multiplicative subgroup of size `domain_size`.
///
/// # Examples
/// ```
/// # use winter_math::fft::*;
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement, StarkField};
/// let n = 2048;
/// let inv_twiddles = get_inv_twiddles::<BaseElement>(n);
///
/// assert_eq!(n / 2, inv_twiddles.len());
/// ```
pub fn get_inv_twiddles<B>(domain_size: usize) -> Vec<B>
where
    B: StarkField,
{
    assert!(
        domain_size.is_power_of_two(),
        "domain size must be a power of 2"
    );
    assert!(
        domain_size.ilog2() <= B::TWO_ADICITY,
        "multiplicative subgroup of size {domain_size} does not exist in the specified base field"
    );
    let root = B::get_root_of_unity(domain_size.ilog2());
    let inv_root = root.exp((domain_size as u32 - 1).into());
    let mut inv_twiddles = get_power_series(inv_root, domain_size / 2);
    permute(&mut inv_twiddles);
    inv_twiddles
}

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

/// Returns the degree of a polynomial implied by the provided evaluations.
///
/// The polynomial is assumed to be evaluated over a multiplicative subgroup of length equal to
/// the length of `evaluations` in the field defined by `B` type parameter and shifted by the
/// `domain_offset`.
///
/// The shifted domain is defined as the original domain with every element multiplied by the
/// `domain_offset`.
///
/// The degree is determined by first interpolating the polynomial into a coefficient form using
/// FFT-based polynomial interpolation, and then finding the first non-zero coefficient.
///
/// # Panics
/// Panics if
/// * Length of `evaluations` is not a power of two.
/// * Field specified by `B` does not contain a multiplicative subgroup of size `domain_size`.
/// * `domain_offset` is ZERO.
///
/// # Examples
/// ```
/// # use winter_math::{polynom, fft::*, get_power_series};
/// # use winter_math::{fields::{f128::BaseElement}, FieldElement, StarkField};
/// let offset = BaseElement::GENERATOR;
/// // p(x) = x^2 + 1
/// let p = [
///     BaseElement::new(1),
///     BaseElement::ZERO,
///     BaseElement::new(1),
///     BaseElement::ZERO,
/// ];
///
/// let g = BaseElement::get_root_of_unity(p.len().ilog2());
/// let domain = get_power_series(g, p.len());
/// let shifted_domain = domain.iter().map(|&x| x * offset).collect::<Vec<_>>();
/// let evaluations = polynom::eval_many(&p, &shifted_domain);
///
/// assert_eq!(2, infer_degree(&evaluations, offset));
/// ```
pub fn infer_degree<B, E>(evaluations: &[E], domain_offset: B) -> usize
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    assert!(
        evaluations.len().is_power_of_two(),
        "number of evaluations must be a power of 2"
    );
    assert!(
        evaluations.len().ilog2() <= B::TWO_ADICITY,
        "multiplicative subgroup of size {} does not exist in the specified base field",
        evaluations.len()
    );
    assert_ne!(domain_offset, B::ZERO, "domain offset cannot be zero");
    let mut poly = evaluations.to_vec();
    let inv_twiddles = get_inv_twiddles::<B>(evaluations.len());
    interpolate_poly_with_offset(&mut poly, &inv_twiddles, domain_offset);
    super::polynom::degree_of(&poly)
}

// PERMUTATIONS
// ================================================================================================

/// Computes bit reverse of the specified index in the domain of the specified size.
///
/// Domain size is assumed to be a power of two and index must be smaller than domain size.
pub fn permute_index(size: usize, index: usize) -> usize {
    const USIZE_BITS: u32 = 0_usize.count_zeros();

    debug_assert!(index < size);
    debug_assert!(size.is_power_of_two());

    let bits = size.trailing_zeros();
    index.reverse_bits().wrapping_shr(USIZE_BITS - bits)
}

fn permute<E: FieldElement>(v: &mut [E]) {
    if cfg!(feature = "concurrent") && v.len() >= MIN_CONCURRENT_SIZE {
        #[cfg(feature = "concurrent")]
        concurrent::permute(v);
    } else {
        FftInputs::permute(v);
    }
}


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

/// Real 2-FFT over u64 integers.
#[inline(always)]
pub fn fft2_real(x: [u64; 2]) -> [i64; 2] {
    [(x[0] as i64 + x[1] as i64), (x[0] as i64 - x[1] as i64)]
}

/// Real 2-iFFT over u64 integers.
/// Division by two to complete the inverse FFT is expected to be performed ***outside*** of this function.
#[inline(always)]
pub fn ifft2_real_unreduced(y: [i64; 2]) -> [u64; 2] {
    [(y[0] + y[1]) as u64, (y[0] - y[1]) as u64]
}

/// Real 4-FFT over u64 integers.
#[inline(always)]
pub fn fft4_real(x: [u64; 4]) -> (i64, (i64, i64), i64) {
    let [z0, z2] = fft2_real([x[0], x[2]]);
    let [z1, z3] = fft2_real([x[1], x[3]]);
    let y0 = z0 + z1;
    let y1 = (z2, -z3);
    let y2 = z0 - z1;
    (y0, y1, y2)
}

/// Real 4-iFFT over u64 integers.
/// Division by four to complete the inverse FFT is expected to be performed ***outside*** of this function.
#[inline(always)]
pub fn ifft4_real_unreduced(y: (i64, (i64, i64), i64)) -> [u64; 4] {
    let z0 = y.0 + y.2;
    let z1 = y.0 - y.2;
    let z2 = y.1 .0;
    let z3 = -y.1 .1;

    let [x0, x2] = ifft2_real_unreduced([z0, z2]);
    let [x1, x3] = ifft2_real_unreduced([z1, z3]);

    [x0, x1, x2, x3]
}


// 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::fft_inputs::FftInputs;
use crate::{field::StarkField, FieldElement};
use utils::{collections::Vec, uninit_vector};

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

/// Evaluates polynomial `p` in-place over the domain of length `p.len()` in the field specified
/// by `B` using the FFT algorithm.
pub fn evaluate_poly<B, E>(p: &mut [E], twiddles: &[B])
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    p.fft_in_place(twiddles);
    p.permute();
}

/// Evaluates polynomial `p` over the domain of length `p.len()` * `blowup_factor` shifted by
/// `domain_offset` in the field specified `B` using the FFT algorithm and returns the result.
pub fn evaluate_poly_with_offset<B, E>(
    p: &[E],
    twiddles: &[B],
    domain_offset: B,
    blowup_factor: usize,
) -> Vec<E>
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    let domain_size = p.len() * blowup_factor;
    let g = B::get_root_of_unity(domain_size.ilog2());
    let mut result = unsafe { uninit_vector(domain_size) };

    result
        .as_mut_slice()
        .chunks_mut(p.len())
        .enumerate()
        .for_each(|(i, chunk)| {
            let idx = super::permute_index(blowup_factor, i) as u64;
            let offset = g.exp(idx.into()) * domain_offset;
            let mut factor = E::BaseField::ONE;
            for (d, c) in chunk.iter_mut().zip(p.iter()) {
                *d = (*c).mul_base(factor);
                factor *= offset;
            }
            chunk.fft_in_place(twiddles);
        });

    result.permute();
    result
}

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

/// Interpolates `evaluations` over a domain of length `evaluations.len()` in the field specified
/// `B` into a polynomial in coefficient form using the FFT algorithm.
pub fn interpolate_poly<B, E>(evaluations: &mut [E], inv_twiddles: &[B])
where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    let inv_length = B::inv((evaluations.len() as u64).into());
    evaluations.fft_in_place(inv_twiddles);
    evaluations.shift_by(inv_length);
    evaluations.permute();
}

/// Interpolates `evaluations` over a domain of length `evaluations.len()` and shifted by
/// `domain_offset` in the field specified by `B` into a polynomial in coefficient form using
/// the FFT algorithm.
pub fn interpolate_poly_with_offset<B, E>(
    evaluations: &mut [E],
    inv_twiddles: &[B],
    domain_offset: B,
) where
    B: StarkField,
    E: FieldElement<BaseField = B>,
{
    evaluations.fft_in_place(inv_twiddles);
    evaluations.permute();

    let domain_offset = B::inv(domain_offset);
    let offset = B::inv((evaluations.len() as u64).into());

    evaluations.shift_by_series(offset, 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.

use crate::{
    fft::fft_inputs::FftInputs,
    field::{f128::BaseElement, StarkField},
    polynom,
    utils::get_power_series,
};
use rand_utils::rand_vector;
use utils::collections::Vec;

// CORE ALGORITHMS
// ================================================================================================

#[test]
fn fft_in_place() {
    // degree 3
    let n = 4;
    let mut p = rand_vector(n);
    let domain = build_domain(n);
    let expected = polynom::eval_many(&p, &domain);
    let twiddles = super::get_twiddles::<BaseElement>(n);
    p.fft_in_place(&twiddles);
    p.permute();
    assert_eq!(expected, p);

    // degree 7
    let n = 8;
    let mut p = rand_vector(n);
    let domain = build_domain(n);
    let twiddles = super::get_twiddles::<BaseElement>(n);
    let expected = polynom::eval_many(&p, &domain);
    p.fft_in_place(&twiddles);
    p.permute();
    assert_eq!(expected, p);

    // degree 15
    let n = 16;
    let mut p = rand_vector(n);
    let domain = build_domain(n);
    let twiddles = super::get_twiddles::<BaseElement>(16);
    let expected = polynom::eval_many(&p, &domain);
    p.fft_in_place(&twiddles);
    p.permute();
    assert_eq!(expected, p);

    // degree 1023
    let n = 1024;
    let mut p = rand_vector(n);
    let domain = build_domain(n);
    let expected = polynom::eval_many(&p, &domain);
    let twiddles = super::get_twiddles::<BaseElement>(n);
    p.fft_in_place(&twiddles);
    p.permute();
    assert_eq!(expected, p);
}

#[test]
fn fft_get_twiddles() {
    let n = super::MIN_CONCURRENT_SIZE * 2;
    let g = BaseElement::get_root_of_unity(n.ilog2());

    let mut expected = get_power_series(g, n / 2);
    expected.permute();

    let twiddles = super::get_twiddles::<BaseElement>(n);
    assert_eq!(expected, twiddles);
}

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

fn build_domain(size: usize) -> Vec<BaseElement> {
    let g = BaseElement::get_root_of_unity(size.ilog2());
    get_power_series(g, size)
}


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

mod traits;
pub use traits::{ExtensibleField, ExtensionOf, FieldElement, StarkField, ToElements};

pub mod f128;
pub mod f62;
pub mod f64;

mod extensions;
pub use extensions::{CubeExtension, QuadExtension};


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

use core::{
    convert::TryFrom,
    fmt::{Debug, Display},
    ops::{
        Add, AddAssign, BitAnd, Div, DivAssign, Mul, MulAssign, Neg, Shl, Shr, ShrAssign, Sub,
        SubAssign,
    },
};
use utils::{
    collections::Vec, AsBytes, Deserializable, DeserializationError, Randomizable, Serializable,
};

// FIELD ELEMENT
// ================================================================================================
/// Defines an element in a finite field.
///
/// This trait defines basic arithmetic operations for elements in
/// [finite fields](https://en.wikipedia.org/wiki/Finite_field) (e.g. addition subtraction,
/// multiplication, division) as well as several convenience functions (e.g. double, square cube).
/// Moreover, it defines interfaces for serializing and deserializing field elements.
///
/// The elements could be in a prime field or an extension of a prime field. Currently, only
/// quadratic and cubic field extensions are supported.
pub trait FieldElement:
    Copy
    + Clone
    + Debug
    + Display
    + Default
    + Send
    + Sync
    + Eq
    + PartialEq
    + Sized
    + Add<Self, Output = Self>
    + Sub<Self, Output = Self>
    + Mul<Self, Output = Self>
    + Div<Self, Output = Self>
    + AddAssign<Self>
    + SubAssign<Self>
    + MulAssign<Self>
    + DivAssign<Self>
    + Neg<Output = Self>
    + From<u128>
    + From<u64>
    + From<u32>
    + From<u16>
    + From<u8>
    + for<'a> TryFrom<&'a [u8]>
    + ExtensionOf<<Self as FieldElement>::BaseField>
    + AsBytes
    + Randomizable
    + Serializable
    + Deserializable
{
    /// A type defining positive integers big enough to describe a field modulus for
    /// `Self::BaseField` with no loss of precision.
    type PositiveInteger: Debug
        + Copy
        + PartialEq
        + PartialOrd
        + ShrAssign
        + Shl<u32, Output = Self::PositiveInteger>
        + Shr<u32, Output = Self::PositiveInteger>
        + BitAnd<Output = Self::PositiveInteger>
        + From<u32>
        + From<u64>;

    /// Base field type for this finite field. For prime fields, `BaseField` should be set
    /// to `Self`.
    type BaseField: StarkField;

    /// Extension degree of this field with respect to `Self::BaseField`. For prime fields,
    /// extension degree should be set to 1.
    const EXTENSION_DEGREE: usize;

    /// Number of bytes needed to encode an element
    const ELEMENT_BYTES: usize;

    /// True if internal representation of the element is the same as its canonical representation.
    const IS_CANONICAL: bool;

    /// The additive identity.
    const ZERO: Self;

    /// The multiplicative identity.
    const ONE: Self;

    // ALGEBRA
    // --------------------------------------------------------------------------------------------

    /// Returns this field element added to itself.
    #[inline]
    #[must_use]
    fn double(self) -> Self {
        self + self
    }

    /// Returns this field element raised to power 2.
    #[inline]
    #[must_use]
    fn square(self) -> Self {
        self * self
    }

    /// Returns this field element raised to power 3.
    #[inline]
    #[must_use]
    fn cube(self) -> Self {
        self * self * self
    }

    /// Exponentiates this field element by `power` parameter.
    #[must_use]
    fn exp(self, power: Self::PositiveInteger) -> Self {
        self.exp_vartime(power)
    }

    /// Exponentiates this field element by `power` parameter.
    /// This function is expressly variable time, to speed-up verifier computations.
    #[must_use]
    fn exp_vartime(self, power: Self::PositiveInteger) -> Self {
        let mut r = Self::ONE;
        let mut b = self;
        let mut p = power;

        let int_zero = Self::PositiveInteger::from(0u32);
        let int_one = Self::PositiveInteger::from(1u32);

        if p == int_zero {
            return Self::ONE;
        } else if b == Self::ZERO {
            return Self::ZERO;
        }

        while p > int_zero {
            if p & int_one == int_one {
                r *= b;
            }
            p >>= int_one;
            b = b.square();
        }

        r
    }

    /// Returns a multiplicative inverse of this field element. If this element is ZERO, ZERO is
    /// returned.
    #[must_use]
    fn inv(self) -> Self;

    /// Returns a conjugate of this field element.
    #[must_use]
    fn conjugate(&self) -> Self;

    // BASE ELEMENT CONVERSIONS
    // --------------------------------------------------------------------------------------------

    /// Return base filed element component of this field element at the specified index `i`.
    ///
    /// # Panics
    /// Panics if the specified index is greater than or equal to `Self::EXTENSION_DEGREE`.
    fn base_element(&self, i: usize) -> Self::BaseField;

    /// Converts a slice of field elements into a slice of elements in the underlying base field.
    ///
    /// For base STARK fields, the input and output slices are the same. For extension fields, the
    /// output slice will contain decompositions of each extension element into underlying base
    /// field elements.
    fn slice_as_base_elements(elements: &[Self]) -> &[Self::BaseField];

    /// Convert a slice of base field elements into a slice of field elements.
    ///
    /// For base STARK fields, the input and output slices are the same. For extension fields, the
    /// output slice will contain a composition of base field elements into extension field
    /// elements.
    ///
    /// # Panics
    /// Panics if the the length of the provided slice is not divisible by `Self::EXTENSION_DEGREE`.
    fn slice_from_base_elements(elements: &[Self::BaseField]) -> &[Self];

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

    /// Converts a list of elements into a list of bytes.
    ///
    /// The elements may be in the internal representation rather than in the canonical
    /// representation. This conversion is intended to be zero-copy (i.e. by re-interpreting the
    /// underlying memory).
    fn elements_as_bytes(elements: &[Self]) -> &[u8];

    /// Converts a list of bytes into a list of field elements.
    ///
    /// The elements are assumed to encoded in the internal representation rather than in the
    /// canonical representation. The conversion is intended to be zero-copy (i.e. by
    /// re-interpreting the underlying memory).
    ///
    /// # Errors
    /// An error is returned if:
    /// * Memory alignment of `bytes` does not match memory alignment of field element data.
    /// * Length of `bytes` does not divide into whole number of elements.
    ///
    /// # Safety
    /// This function is unsafe because it does not check whether underlying bytes represent valid
    /// field elements according to their internal representation.
    unsafe fn bytes_as_elements(bytes: &[u8]) -> Result<&[Self], DeserializationError>;

    // UTILITIES
    // --------------------------------------------------------------------------------------------

    /// Returns a vector of length `n` initialized with all ZERO elements.
    ///
    /// Specialized implementations of this function may be faster than the generic implementation.
    fn zeroed_vector(n: usize) -> Vec<Self> {
        vec![Self::ZERO; n]
    }
}

// STARK FIELD
// ================================================================================================

/// Defines an element in a STARK-friendly finite field.
///
/// A STARK-friendly field is defined as a prime field with high two-addicity. That is, the
/// the modulus of the field should be a prime number of the form `k` * 2^`n` + 1 (a Proth prime),
/// where `n` is relatively large (e.g., greater than 32).
pub trait StarkField: FieldElement<BaseField = Self> {
    /// Prime modulus of the field. Must be of the form `k` * 2^`n` + 1 (a Proth prime).
    /// This ensures that the field has high 2-adicity.
    const MODULUS: Self::PositiveInteger;

    /// The number of bits needed to represents `Self::MODULUS`.
    const MODULUS_BITS: u32;

    /// A multiplicative generator of the field.
    const GENERATOR: Self;

    /// Let Self::MODULUS = `k` * 2^`n` + 1; then, TWO_ADICITY is `n`.
    const TWO_ADICITY: u32;

    /// Let Self::MODULUS = `k` * 2^`n` + 1; then, TWO_ADIC_ROOT_OF_UNITY is 2^`n` root of unity
    /// computed as Self::GENERATOR^`k`.
    const TWO_ADIC_ROOT_OF_UNITY: Self;

    /// Returns the root of unity of order 2^`n`.
    ///
    /// # Panics
    /// Panics if the root of unity for the specified order does not exist in this field.
    fn get_root_of_unity(n: u32) -> Self {
        assert!(n != 0, "cannot get root of unity for n = 0");
        assert!(
            n <= Self::TWO_ADICITY,
            "order cannot exceed 2^{}",
            Self::TWO_ADICITY
        );
        let power = Self::PositiveInteger::from(1u32) << (Self::TWO_ADICITY - n);
        Self::TWO_ADIC_ROOT_OF_UNITY.exp(power)
    }

    /// Returns byte representation of the field modulus in little-endian byte order.
    fn get_modulus_le_bytes() -> Vec<u8>;

    /// Returns a canonical integer representation of this field element.
    fn as_int(&self) -> Self::PositiveInteger;
}

// EXTENSIBLE FIELD
// ================================================================================================

/// Defines basic arithmetic in an extension of a [StarkField] of a given degree.
///
/// This trait defines how to perform multiplication and compute a Frobenius automorphisms of an
/// element in an extension of degree N for a given [StarkField]. It as assumed that an element in
/// degree N extension field can be represented by N field elements in the base field.
///
/// Implementation of this trait implicitly defines the irreducible polynomial over which the
/// extension field is defined.
pub trait ExtensibleField<const N: usize>: StarkField {
    /// Returns a product of `a` and `b` in the field defined by this extension.
    fn mul(a: [Self; N], b: [Self; N]) -> [Self; N];

    /// Returns the square of `a` in the field defined by this extension.
    fn square(a: [Self; N]) -> [Self; N] {
        <Self as ExtensibleField<N>>::mul(a, a)
    }

    /// Returns a product of `a` and `b` in the field defined by this extension. `b` represents
    /// an element in the base field.
    fn mul_base(a: [Self; N], b: Self) -> [Self; N];

    /// Returns Frobenius automorphisms for `x` in the field defined by this extension.
    fn frobenius(x: [Self; N]) -> [Self; N];

    /// Returns true if this extension is supported for the underlying base field.
    fn is_supported() -> bool {
        true
    }
}

// EXTENSION OF
// ================================================================================================

/// Specifies that a field is an extension of another field.
///
/// Currently, this implies the following:
/// - An element in the base field can be converted into an element in the extension field.
/// - An element in the extension field can be multiplied by a base field element directly. This
///   can be used for optimization purposes as such multiplication could be much more efficient
///   than multiplication of two extension field elements.
pub trait ExtensionOf<E: FieldElement>: From<E> {
    fn mul_base(self, other: E) -> Self;
}

/// A field is always an extension of itself.
impl<E: FieldElement> ExtensionOf<E> for E {
    #[inline(always)]
    fn mul_base(self, other: E) -> Self {
        self * other
    }
}

// TO ELEMENTS
// ================================================================================================

/// Defines how to convert a struct to a vector of field elements.
pub trait ToElements<E: FieldElement> {
    fn to_elements(&self) -> Vec<E>;
}

impl<E: FieldElement> ToElements<E> for () {
    fn to_elements(&self) -> Vec<E> {
        Vec::new()
    }
}

impl<E: FieldElement> ToElements<E> for E {
    fn to_elements(&self) -> Vec<E> {
        vec![*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 super::{ExtensibleField, ExtensionOf, FieldElement};
use core::{
    convert::TryFrom,
    fmt,
    ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign},
    slice,
};
use utils::{
    collections::Vec, string::ToString, AsBytes, ByteReader, ByteWriter, Deserializable,
    DeserializationError, Randomizable, Serializable, SliceReader,
};

// QUADRATIC EXTENSION FIELD
// ================================================================================================

/// Represents an element in a cubic extension of a [StarkField](crate::StarkField).
///
/// The extension element is defined as Î± + Î² * Ï† + Î³ * Ï†^2, where Ï† is a root of in irreducible
/// polynomial defined by the implementation of the [ExtensibleField] trait, and Î±, Î², Î³ are base
/// field elements.
#[repr(C)]
#[derive(Copy, Clone, Debug, PartialEq, Eq, Default)]
pub struct CubeExtension<B: ExtensibleField<3>>(B, B, B);

impl<B: ExtensibleField<3>> CubeExtension<B> {
    /// Returns a new extension element instantiated from the provided base elements.
    pub const fn new(a: B, b: B, c: B) -> Self {
        Self(a, b, c)
    }

    /// Returns true if the base field specified by B type parameter supports cubic extensions.
    pub fn is_supported() -> bool {
        <B as ExtensibleField<3>>::is_supported()
    }

    /// Converts a vector of base elements into a vector of elements in a cubic extension field
    /// by fusing three adjacent base elements together. The output vector is one-third the length
    /// of the source vector.
    fn base_to_cubic_vector(source: Vec<B>) -> Vec<Self> {
        debug_assert!(
            source.len() % Self::EXTENSION_DEGREE == 0,
            "source vector length must be divisible by three, but was {}",
            source.len()
        );
        let mut v = core::mem::ManuallyDrop::new(source);
        let p = v.as_mut_ptr();
        let len = v.len() / Self::EXTENSION_DEGREE;
        let cap = v.capacity() / Self::EXTENSION_DEGREE;
        unsafe { Vec::from_raw_parts(p as *mut Self, len, cap) }
    }

    /// Returns an array of base field elements comprising this extension field element.
    ///
    /// The order of abase elements in the returned array is the same as the order in which
    /// the elements are provided to the [CubeExtension::new()] constructor.
    pub const fn to_base_elements(self) -> [B; 3] {
        [self.0, self.1, self.2]
    }
}

impl<B: ExtensibleField<3>> FieldElement for CubeExtension<B> {
    type PositiveInteger = B::PositiveInteger;
    type BaseField = B;

    const EXTENSION_DEGREE: usize = 3;

    const ELEMENT_BYTES: usize = B::ELEMENT_BYTES * Self::EXTENSION_DEGREE;
    const IS_CANONICAL: bool = B::IS_CANONICAL;
    const ZERO: Self = Self(B::ZERO, B::ZERO, B::ZERO);
    const ONE: Self = Self(B::ONE, B::ZERO, B::ZERO);

    // ALGEBRA
    // --------------------------------------------------------------------------------------------

    #[inline]
    fn double(self) -> Self {
        Self(self.0.double(), self.1.double(), self.2.double())
    }

    #[inline]
    fn square(self) -> Self {
        let a = <B as ExtensibleField<3>>::square([self.0, self.1, self.2]);
        Self(a[0], a[1], a[2])
    }

    #[inline]
    fn inv(self) -> Self {
        if self == Self::ZERO {
            return self;
        }

        let x = [self.0, self.1, self.2];
        let c1 = <B as ExtensibleField<3>>::frobenius(x);
        let c2 = <B as ExtensibleField<3>>::frobenius(c1);
        let numerator = <B as ExtensibleField<3>>::mul(c1, c2);

        let norm = <B as ExtensibleField<3>>::mul(x, numerator);
        debug_assert_eq!(norm[1], B::ZERO, "norm must be in the base field");
        debug_assert_eq!(norm[2], B::ZERO, "norm must be in the base field");
        let denom_inv = norm[0].inv();

        Self(
            numerator[0] * denom_inv,
            numerator[1] * denom_inv,
            numerator[2] * denom_inv,
        )
    }

    #[inline]
    fn conjugate(&self) -> Self {
        let result = <B as ExtensibleField<3>>::frobenius([self.0, self.1, self.2]);
        Self(result[0], result[1], result[2])
    }

    // BASE ELEMENT CONVERSIONS
    // --------------------------------------------------------------------------------------------

    fn base_element(&self, i: usize) -> Self::BaseField {
        match i {
            0 => self.0,
            1 => self.1,
            2 => self.2,
            _ => panic!("element index must be smaller than 3, but was {i}"),
        }
    }

    fn slice_as_base_elements(elements: &[Self]) -> &[Self::BaseField] {
        let ptr = elements.as_ptr();
        let len = elements.len() * Self::EXTENSION_DEGREE;
        unsafe { slice::from_raw_parts(ptr as *const Self::BaseField, len) }
    }

    fn slice_from_base_elements(elements: &[Self::BaseField]) -> &[Self] {
        assert!(
            elements.len() % Self::EXTENSION_DEGREE == 0,
            "number of base elements must be divisible by 3, but was {}",
            elements.len()
        );

        let ptr = elements.as_ptr();
        let len = elements.len() / Self::EXTENSION_DEGREE;
        unsafe { slice::from_raw_parts(ptr as *const Self, len) }
    }

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

    fn elements_as_bytes(elements: &[Self]) -> &[u8] {
        unsafe {
            slice::from_raw_parts(
                elements.as_ptr() as *const u8,
                elements.len() * Self::ELEMENT_BYTES,
            )
        }
    }

    unsafe fn bytes_as_elements(bytes: &[u8]) -> Result<&[Self], DeserializationError> {
        if bytes.len() % Self::ELEMENT_BYTES != 0 {
            return Err(DeserializationError::InvalidValue(format!(
                "number of bytes ({}) does not divide into whole number of field elements",
                bytes.len(),
            )));
        }

        let p = bytes.as_ptr();
        let len = bytes.len() / Self::ELEMENT_BYTES;

        // make sure the bytes are aligned on the boundary consistent with base element alignment
        if (p as usize) % Self::BaseField::ELEMENT_BYTES != 0 {
            return Err(DeserializationError::InvalidValue(
                "slice memory alignment is not valid for this field element type".to_string(),
            ));
        }

        Ok(slice::from_raw_parts(p as *const Self, len))
    }

    // UTILITIES
    // --------------------------------------------------------------------------------------------

    fn zeroed_vector(n: usize) -> Vec<Self> {
        // get three times the number of base elements and re-interpret them as cubic field
        // elements
        let result = B::zeroed_vector(n * Self::EXTENSION_DEGREE);
        Self::base_to_cubic_vector(result)
    }
}

impl<B: ExtensibleField<3>> ExtensionOf<B> for CubeExtension<B> {
    #[inline(always)]
    fn mul_base(self, other: B) -> Self {
        let result = <B as ExtensibleField<3>>::mul_base([self.0, self.1, self.2], other);
        Self(result[0], result[1], result[2])
    }
}

impl<B: ExtensibleField<3>> Randomizable for CubeExtension<B> {
    const VALUE_SIZE: usize = Self::ELEMENT_BYTES;

    fn from_random_bytes(bytes: &[u8]) -> Option<Self> {
        Self::try_from(bytes).ok()
    }
}

impl<B: ExtensibleField<3>> fmt::Display for CubeExtension<B> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "({}, {}, {})", self.0, self.1, self.2)
    }
}

// OVERLOADED OPERATORS
// ------------------------------------------------------------------------------------------------

impl<B: ExtensibleField<3>> Add for CubeExtension<B> {
    type Output = Self;

    #[inline]
    fn add(self, rhs: Self) -> Self {
        Self(self.0 + rhs.0, self.1 + rhs.1, self.2 + rhs.2)
    }
}

impl<B: ExtensibleField<3>> AddAssign for CubeExtension<B> {
    #[inline]
    fn add_assign(&mut self, rhs: Self) {
        *self = *self + rhs
    }
}

impl<B: ExtensibleField<3>> Sub for CubeExtension<B> {
    type Output = Self;

    #[inline]
    fn sub(self, rhs: Self) -> Self {
        Self(self.0 - rhs.0, self.1 - rhs.1, self.2 - rhs.2)
    }
}

impl<B: ExtensibleField<3>> SubAssign for CubeExtension<B> {
    #[inline]
    fn sub_assign(&mut self, rhs: Self) {
        *self = *self - rhs;
    }
}

impl<B: ExtensibleField<3>> Mul for CubeExtension<B> {
    type Output = Self;

    #[inline]
    fn mul(self, rhs: Self) -> Self {
        let result =
            <B as ExtensibleField<3>>::mul([self.0, self.1, self.2], [rhs.0, rhs.1, rhs.2]);
        Self(result[0], result[1], result[2])
    }
}

impl<B: ExtensibleField<3>> MulAssign for CubeExtension<B> {
    #[inline]
    fn mul_assign(&mut self, rhs: Self) {
        *self = *self * rhs
    }
}

impl<B: ExtensibleField<3>> Div for CubeExtension<B> {
    type Output = Self;

    #[inline]
    #[allow(clippy::suspicious_arithmetic_impl)]
    fn div(self, rhs: Self) -> Self {
        self * rhs.inv()
    }
}

impl<B: ExtensibleField<3>> DivAssign for CubeExtension<B> {
    #[inline]
    fn div_assign(&mut self, rhs: Self) {
        *self = *self / rhs
    }
}

impl<B: ExtensibleField<3>> Neg for CubeExtension<B> {
    type Output = Self;

    #[inline]
    fn neg(self) -> Self {
        Self(-self.0, -self.1, -self.2)
    }
}

// TYPE CONVERSIONS
// ------------------------------------------------------------------------------------------------

impl<B: ExtensibleField<3>> From<B> for CubeExtension<B> {
    fn from(value: B) -> Self {
        Self(value, B::ZERO, B::ZERO)
    }
}

impl<B: ExtensibleField<3>> From<u128> for CubeExtension<B> {
    fn from(value: u128) -> Self {
        Self(B::from(value), B::ZERO, B::ZERO)
    }
}

impl<B: ExtensibleField<3>> From<u64> for CubeExtension<B> {
    fn from(value: u64) -> Self {
        Self(B::from(value), B::ZERO, B::ZERO)
    }
}

impl<B: ExtensibleField<3>> From<u32> for CubeExtension<B> {
    fn from(value: u32) -> Self {
        Self(B::from(value), B::ZERO, B::ZERO)
    }
}

impl<B: ExtensibleField<3>> From<u16> for CubeExtension<B> {
    fn from(value: u16) -> Self {
        Self(B::from(value), B::ZERO, B::ZERO)
    }
}

impl<B: ExtensibleField<3>> From<u8> for CubeExtension<B> {
    fn from(value: u8) -> Self {
        Self(B::from(value), B::ZERO, B::ZERO)
    }
}

impl<'a, B: ExtensibleField<3>> TryFrom<&'a [u8]> for CubeExtension<B> {
    type Error = DeserializationError;

    /// Converts a slice of bytes into a field element; returns error if the value encoded in bytes
    /// is not a valid field element. The bytes are assumed to be in little-endian byte order.
    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
        if bytes.len() < Self::ELEMENT_BYTES {
            return Err(DeserializationError::InvalidValue(format!(
                "not enough bytes for a full field element; expected {} bytes, but was {} bytes",
                Self::ELEMENT_BYTES,
                bytes.len(),
            )));
        }
        if bytes.len() > Self::ELEMENT_BYTES {
            return Err(DeserializationError::InvalidValue(format!(
                "too many bytes for a field element; expected {} bytes, but was {} bytes",
                Self::ELEMENT_BYTES,
                bytes.len(),
            )));
        }
        let mut reader = SliceReader::new(bytes);
        Self::read_from(&mut reader)
    }
}

impl<B: ExtensibleField<3>> AsBytes for CubeExtension<B> {
    fn as_bytes(&self) -> &[u8] {
        // TODO: take endianness into account
        let self_ptr: *const Self = self;
        unsafe { slice::from_raw_parts(self_ptr as *const u8, Self::ELEMENT_BYTES) }
    }
}

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

impl<B: ExtensibleField<3>> Serializable for CubeExtension<B> {
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        self.0.write_into(target);
        self.1.write_into(target);
        self.2.write_into(target);
    }
}

impl<B: ExtensibleField<3>> Deserializable for CubeExtension<B> {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let value0 = B::read_from(source)?;
        let value1 = B::read_from(source)?;
        let value2 = B::read_from(source)?;
        Ok(Self(value0, value1, value2))
    }
}

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

#[cfg(test)]
mod tests {
    use super::{CubeExtension, DeserializationError, FieldElement};
    use crate::field::f64::BaseElement;
    use rand_utils::rand_value;

    // BASIC ALGEBRA
    // --------------------------------------------------------------------------------------------

    #[test]
    fn add() {
        // identity
        let r: CubeExtension<BaseElement> = rand_value();
        assert_eq!(r, r + CubeExtension::<BaseElement>::ZERO);

        // test random values
        let r1: CubeExtension<BaseElement> = rand_value();
        let r2: CubeExtension<BaseElement> = rand_value();

        let expected = CubeExtension(r1.0 + r2.0, r1.1 + r2.1, r1.2 + r2.2);
        assert_eq!(expected, r1 + r2);
    }

    #[test]
    fn sub() {
        // identity
        let r: CubeExtension<BaseElement> = rand_value();
        assert_eq!(r, r - CubeExtension::<BaseElement>::ZERO);

        // test random values
        let r1: CubeExtension<BaseElement> = rand_value();
        let r2: CubeExtension<BaseElement> = rand_value();

        let expected = CubeExtension(r1.0 - r2.0, r1.1 - r2.1, r1.2 - r2.2);
        assert_eq!(expected, r1 - r2);
    }

    // INITIALIZATION
    // --------------------------------------------------------------------------------------------

    #[test]
    fn zeroed_vector() {
        let result = CubeExtension::<BaseElement>::zeroed_vector(4);
        assert_eq!(4, result.len());
        for element in result.into_iter() {
            assert_eq!(CubeExtension::<BaseElement>::ZERO, element);
        }
    }

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

    #[test]
    fn elements_as_bytes() {
        let source = vec![
            CubeExtension(
                BaseElement::new(1),
                BaseElement::new(2),
                BaseElement::new(3),
            ),
            CubeExtension(
                BaseElement::new(4),
                BaseElement::new(5),
                BaseElement::new(6),
            ),
        ];

        let mut expected = vec![];
        expected.extend_from_slice(&source[0].0.inner().to_le_bytes());
        expected.extend_from_slice(&source[0].1.inner().to_le_bytes());
        expected.extend_from_slice(&source[0].2.inner().to_le_bytes());
        expected.extend_from_slice(&source[1].0.inner().to_le_bytes());
        expected.extend_from_slice(&source[1].1.inner().to_le_bytes());
        expected.extend_from_slice(&source[1].2.inner().to_le_bytes());

        assert_eq!(
            expected,
            CubeExtension::<BaseElement>::elements_as_bytes(&source)
        );
    }

    #[test]
    fn bytes_as_elements() {
        let elements = vec![
            CubeExtension(
                BaseElement::new(1),
                BaseElement::new(2),
                BaseElement::new(3),
            ),
            CubeExtension(
                BaseElement::new(4),
                BaseElement::new(5),
                BaseElement::new(6),
            ),
        ];

        let mut bytes = vec![];
        bytes.extend_from_slice(&elements[0].0.inner().to_le_bytes());
        bytes.extend_from_slice(&elements[0].1.inner().to_le_bytes());
        bytes.extend_from_slice(&elements[0].2.inner().to_le_bytes());
        bytes.extend_from_slice(&elements[1].0.inner().to_le_bytes());
        bytes.extend_from_slice(&elements[1].1.inner().to_le_bytes());
        bytes.extend_from_slice(&elements[1].2.inner().to_le_bytes());
        bytes.extend_from_slice(&BaseElement::new(5).inner().to_le_bytes());

        let result = unsafe { CubeExtension::<BaseElement>::bytes_as_elements(&bytes[..48]) };
        assert!(result.is_ok());
        assert_eq!(elements, result.unwrap());

        let result = unsafe { CubeExtension::<BaseElement>::bytes_as_elements(&bytes) };
        assert!(matches!(result, Err(DeserializationError::InvalidValue(_))));

        let result = unsafe { CubeExtension::<BaseElement>::bytes_as_elements(&bytes[1..]) };
        assert!(matches!(result, Err(DeserializationError::InvalidValue(_))));
    }

    // UTILITIES
    // --------------------------------------------------------------------------------------------

    #[test]
    fn as_base_elements() {
        let elements = vec![
            CubeExtension(
                BaseElement::new(1),
                BaseElement::new(2),
                BaseElement::new(3),
            ),
            CubeExtension(
                BaseElement::new(4),
                BaseElement::new(5),
                BaseElement::new(6),
            ),
        ];

        let expected = vec![
            BaseElement::new(1),
            BaseElement::new(2),
            BaseElement::new(3),
            BaseElement::new(4),
            BaseElement::new(5),
            BaseElement::new(6),
        ];

        assert_eq!(
            expected,
            CubeExtension::<BaseElement>::slice_as_base_elements(&elements)
        );
    }
}


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

mod quadratic;
pub use quadratic::QuadExtension;

mod cubic;
pub use cubic::CubeExtension;

use super::{ExtensibleField, ExtensionOf, FieldElement};


// 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::{ExtensibleField, ExtensionOf, FieldElement};
use core::{
    convert::TryFrom,
    fmt,
    ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign},
    slice,
};
use utils::{
    collections::Vec, string::ToString, AsBytes, ByteReader, ByteWriter, Deserializable,
    DeserializationError, Randomizable, Serializable, SliceReader,
};

// QUADRATIC EXTENSION FIELD
// ================================================================================================

/// Represents an element in a quadratic extension of a [StarkField](crate::StarkField).
///
/// The extension element is defined as Î± + Î² * Ï†, where Ï† is a root of in irreducible polynomial
/// defined by the implementation of the [ExtensibleField] trait, and Î± and Î² are base field
/// elements.
#[repr(C)]
#[derive(Copy, Clone, Debug, PartialEq, Eq, Default)]
pub struct QuadExtension<B: ExtensibleField<2>>(B, B);

impl<B: ExtensibleField<2>> QuadExtension<B> {
    /// Returns a new extension element instantiated from the provided base elements.
    pub const fn new(a: B, b: B) -> Self {
        Self(a, b)
    }

    /// Returns true if the base field specified by B type parameter supports quadratic extensions.
    pub fn is_supported() -> bool {
        <B as ExtensibleField<2>>::is_supported()
    }

    /// Converts a vector of base elements into a vector of elements in a quadratic extension
    /// field by fusing two adjacent base elements together. The output vector is half the length
    /// of the source vector.
    fn base_to_quad_vector(source: Vec<B>) -> Vec<Self> {
        debug_assert!(
            source.len() % Self::EXTENSION_DEGREE == 0,
            "source vector length must be divisible by two, but was {}",
            source.len()
        );
        let mut v = core::mem::ManuallyDrop::new(source);
        let p = v.as_mut_ptr();
        let len = v.len() / Self::EXTENSION_DEGREE;
        let cap = v.capacity() / Self::EXTENSION_DEGREE;
        unsafe { Vec::from_raw_parts(p as *mut Self, len, cap) }
    }

    /// Returns an array of base field elements comprising this extension field element.
    ///
    /// The order of abase elements in the returned array is the same as the order in which
    /// the elements are provided to the [QuadExtension::new()] constructor.
    pub const fn to_base_elements(self) -> [B; 2] {
        [self.0, self.1]
    }
}

impl<B: ExtensibleField<2>> FieldElement for QuadExtension<B> {
    type PositiveInteger = B::PositiveInteger;
    type BaseField = B;

    const EXTENSION_DEGREE: usize = 2;

    const ELEMENT_BYTES: usize = B::ELEMENT_BYTES * Self::EXTENSION_DEGREE;
    const IS_CANONICAL: bool = B::IS_CANONICAL;
    const ZERO: Self = Self(B::ZERO, B::ZERO);
    const ONE: Self = Self(B::ONE, B::ZERO);

    // ALGEBRA
    // --------------------------------------------------------------------------------------------

    #[inline]
    fn double(self) -> Self {
        Self(self.0.double(), self.1.double())
    }

    #[inline]
    fn square(self) -> Self {
        let a = <B as ExtensibleField<2>>::square([self.0, self.1]);
        Self(a[0], a[1])
    }

    #[inline]
    fn inv(self) -> Self {
        if self == Self::ZERO {
            return self;
        }

        let x = [self.0, self.1];
        let numerator = <B as ExtensibleField<2>>::frobenius(x);

        let norm = <B as ExtensibleField<2>>::mul(x, numerator);
        debug_assert_eq!(norm[1], B::ZERO, "norm must be in the base field");
        let denom_inv = norm[0].inv();

        Self(numerator[0] * denom_inv, numerator[1] * denom_inv)
    }

    #[inline]
    fn conjugate(&self) -> Self {
        let result = <B as ExtensibleField<2>>::frobenius([self.0, self.1]);
        Self(result[0], result[1])
    }

    // BASE ELEMENT CONVERSIONS
    // --------------------------------------------------------------------------------------------

    fn base_element(&self, i: usize) -> Self::BaseField {
        match i {
            0 => self.0,
            1 => self.1,
            _ => panic!("element index must be smaller than 2, but was {i}"),
        }
    }

    fn slice_as_base_elements(elements: &[Self]) -> &[Self::BaseField] {
        let ptr = elements.as_ptr();
        let len = elements.len() * Self::EXTENSION_DEGREE;
        unsafe { slice::from_raw_parts(ptr as *const Self::BaseField, len) }
    }

    fn slice_from_base_elements(elements: &[Self::BaseField]) -> &[Self] {
        assert!(
            elements.len() % Self::EXTENSION_DEGREE == 0,
            "number of base elements must be divisible by 2, but was {}",
            elements.len()
        );

        let ptr = elements.as_ptr();
        let len = elements.len() / Self::EXTENSION_DEGREE;
        unsafe { slice::from_raw_parts(ptr as *const Self, len) }
    }

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

    fn elements_as_bytes(elements: &[Self]) -> &[u8] {
        unsafe {
            slice::from_raw_parts(
                elements.as_ptr() as *const u8,
                elements.len() * Self::ELEMENT_BYTES,
            )
        }
    }

    unsafe fn bytes_as_elements(bytes: &[u8]) -> Result<&[Self], DeserializationError> {
        if bytes.len() % Self::ELEMENT_BYTES != 0 {
            return Err(DeserializationError::InvalidValue(format!(
                "number of bytes ({}) does not divide into whole number of field elements",
                bytes.len(),
            )));
        }

        let p = bytes.as_ptr();
        let len = bytes.len() / Self::ELEMENT_BYTES;

        // make sure the bytes are aligned on the boundary consistent with base element alignment
        if (p as usize) % Self::BaseField::ELEMENT_BYTES != 0 {
            return Err(DeserializationError::InvalidValue(
                "slice memory alignment is not valid for this field element type".to_string(),
            ));
        }

        Ok(slice::from_raw_parts(p as *const Self, len))
    }

    // UTILITIES
    // --------------------------------------------------------------------------------------------

    fn zeroed_vector(n: usize) -> Vec<Self> {
        // get twice the number of base elements, and re-interpret them as quad field elements
        let result = B::zeroed_vector(n * Self::EXTENSION_DEGREE);
        Self::base_to_quad_vector(result)
    }
}

impl<B: ExtensibleField<2>> ExtensionOf<B> for QuadExtension<B> {
    #[inline(always)]
    fn mul_base(self, other: B) -> Self {
        let result = <B as ExtensibleField<2>>::mul_base([self.0, self.1], other);
        Self(result[0], result[1])
    }
}

impl<B: ExtensibleField<2>> Randomizable for QuadExtension<B> {
    const VALUE_SIZE: usize = Self::ELEMENT_BYTES;

    fn from_random_bytes(bytes: &[u8]) -> Option<Self> {
        Self::try_from(bytes).ok()
    }
}

impl<B: ExtensibleField<2>> fmt::Display for QuadExtension<B> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "({}, {})", self.0, self.1)
    }
}

// OVERLOADED OPERATORS
// ------------------------------------------------------------------------------------------------

impl<B: ExtensibleField<2>> Add for QuadExtension<B> {
    type Output = Self;

    #[inline]
    fn add(self, rhs: Self) -> Self {
        Self(self.0 + rhs.0, self.1 + rhs.1)
    }
}

impl<B: ExtensibleField<2>> AddAssign for QuadExtension<B> {
    #[inline]
    fn add_assign(&mut self, rhs: Self) {
        *self = *self + rhs
    }
}

impl<B: ExtensibleField<2>> Sub for QuadExtension<B> {
    type Output = Self;

    #[inline]
    fn sub(self, rhs: Self) -> Self {
        Self(self.0 - rhs.0, self.1 - rhs.1)
    }
}

impl<B: ExtensibleField<2>> SubAssign for QuadExtension<B> {
    #[inline]
    fn sub_assign(&mut self, rhs: Self) {
        *self = *self - rhs;
    }
}

impl<B: ExtensibleField<2>> Mul for QuadExtension<B> {
    type Output = Self;

    #[inline]
    fn mul(self, rhs: Self) -> Self {
        let result = <B as ExtensibleField<2>>::mul([self.0, self.1], [rhs.0, rhs.1]);
        Self(result[0], result[1])
    }
}

impl<B: ExtensibleField<2>> MulAssign for QuadExtension<B> {
    #[inline]
    fn mul_assign(&mut self, rhs: Self) {
        *self = *self * rhs
    }
}

impl<B: ExtensibleField<2>> Div for QuadExtension<B> {
    type Output = Self;

    #[inline]
    #[allow(clippy::suspicious_arithmetic_impl)]
    fn div(self, rhs: Self) -> Self {
        self * rhs.inv()
    }
}

impl<B: ExtensibleField<2>> DivAssign for QuadExtension<B> {
    #[inline]
    fn div_assign(&mut self, rhs: Self) {
        *self = *self / rhs
    }
}

impl<B: ExtensibleField<2>> Neg for QuadExtension<B> {
    type Output = Self;

    #[inline]
    fn neg(self) -> Self {
        Self(-self.0, -self.1)
    }
}

// TYPE CONVERSIONS
// ------------------------------------------------------------------------------------------------

impl<B: ExtensibleField<2>> From<B> for QuadExtension<B> {
    fn from(value: B) -> Self {
        Self(value, B::ZERO)
    }
}

impl<B: ExtensibleField<2>> From<u128> for QuadExtension<B> {
    fn from(value: u128) -> Self {
        Self(B::from(value), B::ZERO)
    }
}

impl<B: ExtensibleField<2>> From<u64> for QuadExtension<B> {
    fn from(value: u64) -> Self {
        Self(B::from(value), B::ZERO)
    }
}

impl<B: ExtensibleField<2>> From<u32> for QuadExtension<B> {
    fn from(value: u32) -> Self {
        Self(B::from(value), B::ZERO)
    }
}

impl<B: ExtensibleField<2>> From<u16> for QuadExtension<B> {
    fn from(value: u16) -> Self {
        Self(B::from(value), B::ZERO)
    }
}

impl<B: ExtensibleField<2>> From<u8> for QuadExtension<B> {
    fn from(value: u8) -> Self {
        Self(B::from(value), B::ZERO)
    }
}

impl<'a, B: ExtensibleField<2>> TryFrom<&'a [u8]> for QuadExtension<B> {
    type Error = DeserializationError;

    /// Converts a slice of bytes into a field element; returns error if the value encoded in bytes
    /// is not a valid field element. The bytes are assumed to be in little-endian byte order.
    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
        if bytes.len() < Self::ELEMENT_BYTES {
            return Err(DeserializationError::InvalidValue(format!(
                "not enough bytes for a full field element; expected {} bytes, but was {} bytes",
                Self::ELEMENT_BYTES,
                bytes.len(),
            )));
        }
        if bytes.len() > Self::ELEMENT_BYTES {
            return Err(DeserializationError::InvalidValue(format!(
                "too many bytes for a field element; expected {} bytes, but was {} bytes",
                Self::ELEMENT_BYTES,
                bytes.len(),
            )));
        }
        let mut reader = SliceReader::new(bytes);
        Self::read_from(&mut reader)
    }
}

impl<B: ExtensibleField<2>> AsBytes for QuadExtension<B> {
    fn as_bytes(&self) -> &[u8] {
        // TODO: take endianness into account
        let self_ptr: *const Self = self;
        unsafe { slice::from_raw_parts(self_ptr as *const u8, Self::ELEMENT_BYTES) }
    }
}

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

impl<B: ExtensibleField<2>> Serializable for QuadExtension<B> {
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        self.0.write_into(target);
        self.1.write_into(target);
    }
}

impl<B: ExtensibleField<2>> Deserializable for QuadExtension<B> {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let value0 = B::read_from(source)?;
        let value1 = B::read_from(source)?;
        Ok(Self(value0, value1))
    }
}

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

#[cfg(test)]
mod tests {
    use super::{DeserializationError, FieldElement, QuadExtension};
    use crate::field::f64::BaseElement;
    use rand_utils::rand_value;

    // BASIC ALGEBRA
    // --------------------------------------------------------------------------------------------

    #[test]
    fn add() {
        // identity
        let r: QuadExtension<BaseElement> = rand_value();
        assert_eq!(r, r + QuadExtension::<BaseElement>::ZERO);

        // test random values
        let r1: QuadExtension<BaseElement> = rand_value();
        let r2: QuadExtension<BaseElement> = rand_value();

        let expected = QuadExtension(r1.0 + r2.0, r1.1 + r2.1);
        assert_eq!(expected, r1 + r2);
    }

    #[test]
    fn sub() {
        // identity
        let r: QuadExtension<BaseElement> = rand_value();
        assert_eq!(r, r - QuadExtension::<BaseElement>::ZERO);

        // test random values
        let r1: QuadExtension<BaseElement> = rand_value();
        let r2: QuadExtension<BaseElement> = rand_value();

        let expected = QuadExtension(r1.0 - r2.0, r1.1 - r2.1);
        assert_eq!(expected, r1 - r2);
    }

    // INITIALIZATION
    // --------------------------------------------------------------------------------------------

    #[test]
    fn zeroed_vector() {
        let result = QuadExtension::<BaseElement>::zeroed_vector(4);
        assert_eq!(4, result.len());
        for element in result.into_iter() {
            assert_eq!(QuadExtension::<BaseElement>::ZERO, element);
        }
    }

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

    #[test]
    fn elements_as_bytes() {
        let source = vec![
            QuadExtension(BaseElement::new(1), BaseElement::new(2)),
            QuadExtension(BaseElement::new(3), BaseElement::new(4)),
        ];

        let mut expected = vec![];
        expected.extend_from_slice(&source[0].0.inner().to_le_bytes());
        expected.extend_from_slice(&source[0].1.inner().to_le_bytes());
        expected.extend_from_slice(&source[1].0.inner().to_le_bytes());
        expected.extend_from_slice(&source[1].1.inner().to_le_bytes());

        assert_eq!(
            expected,
            QuadExtension::<BaseElement>::elements_as_bytes(&source)
        );
    }

    #[test]
    fn bytes_as_elements() {
        let elements = vec![
            QuadExtension(BaseElement::new(1), BaseElement::new(2)),
            QuadExtension(BaseElement::new(3), BaseElement::new(4)),
        ];

        let mut bytes = vec![];
        bytes.extend_from_slice(&elements[0].0.inner().to_le_bytes());
        bytes.extend_from_slice(&elements[0].1.inner().to_le_bytes());
        bytes.extend_from_slice(&elements[1].0.inner().to_le_bytes());
        bytes.extend_from_slice(&elements[1].1.inner().to_le_bytes());
        bytes.extend_from_slice(&BaseElement::new(5).inner().to_le_bytes());
        let result = unsafe { QuadExtension::<BaseElement>::bytes_as_elements(&bytes[..32]) };
        assert!(result.is_ok());
        assert_eq!(elements, result.unwrap());

        let result = unsafe { QuadExtension::<BaseElement>::bytes_as_elements(&bytes) };
        assert!(matches!(result, Err(DeserializationError::InvalidValue(_))));

        let result = unsafe { QuadExtension::<BaseElement>::bytes_as_elements(&bytes[1..]) };
        assert!(matches!(result, Err(DeserializationError::InvalidValue(_))));
    }

    // UTILITIES
    // --------------------------------------------------------------------------------------------

    #[test]
    fn as_base_elements() {
        let elements = vec![
            QuadExtension(BaseElement::new(1), BaseElement::new(2)),
            QuadExtension(BaseElement::new(3), BaseElement::new(4)),
        ];

        let expected = vec![
            BaseElement::new(1),
            BaseElement::new(2),
            BaseElement::new(3),
            BaseElement::new(4),
        ];

        assert_eq!(
            expected,
            QuadExtension::<BaseElement>::slice_as_base_elements(&elements)
        );
    }
}