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

use super::ColMatrix;
use core::ops::Deref;
use math::{fft::fft_inputs::FftInputs, FieldElement, StarkField};
use utils::{collections::Vec, group_vector_elements, uninit_vector};

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

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

/// Segments with domain sizes under this number will be evaluated in a single thread.
const MIN_CONCURRENT_SIZE: usize = 1024;

// SEGMENT OF ROW-MAJOR MATRIX
// ================================================================================================

/// A set of columns of a matrix stored in row-major form.
///
/// The rows are stored in a single vector where each element is an array of size `N`. A segment
/// can store [StarkField] elements only, but can be instantiated from a [Matrix] of any extension
/// of the specified [StarkField]. In such a case, extension field elements are decomposed into
/// base field elements and then added to the segment.
#[derive(Clone, Debug)]
pub struct Segment<B: StarkField, const N: usize> {
    data: Vec<[B; N]>,
}

impl<B: StarkField, const N: usize> Segment<B, N> {
    // CONSTRUCTORS
    // --------------------------------------------------------------------------------------------

    /// Instantiates a new [Segment] by evaluating polynomials from the provided [ColMatrix]
    /// starting at the specified offset.
    ///
    /// The offset is assumed to be an offset into the view of the matrix where extension field
    /// elements are decomposed into base field elements. This offset must be compatible with the
    /// values supplied into [Matrix::get_base_element()] method.
    ///
    /// Evaluation is performed over the domain specified by the provided twiddles and offsets.
    ///
    /// # Panics
    /// Panics if:
    /// - `poly_offset` greater than or equal to the number of base field columns in `polys`.
    /// - Number of offsets is not a power of two.
    /// - Number of offsets is smaller than or equal to the polynomial size.
    /// - The number of twiddles is not half the size of the polynomial size.
    pub fn new<E>(polys: &ColMatrix<E>, poly_offset: usize, offsets: &[B], twiddles: &[B]) -> Self
    where
        E: FieldElement<BaseField = B>,
    {
        let poly_size = polys.num_rows();
        let domain_size = offsets.len();
        assert!(domain_size.is_power_of_two());
        assert!(domain_size > poly_size);
        assert_eq!(poly_size, twiddles.len() * 2);
        assert!(poly_offset < polys.num_base_cols());

        // allocate memory for the segment
        let data = if polys.num_base_cols() - poly_offset >= N {
            // if we will fill the entire segment, we allocate uninitialized memory
            unsafe { uninit_vector::<[B; N]>(domain_size) }
        } else {
            // but if some columns in the segment will remain unfilled, we allocate memory initialized
            // to zeros to make sure we don't end up with memory with undefined values
            group_vector_elements(B::zeroed_vector(N * domain_size))
        };

        Self::new_with_buffer(data, polys, poly_offset, offsets, twiddles)
    }

    /// Instantiates a new [Segment] using the provided data buffer by evaluating polynomials in
    /// the [ColMatrix] starting at the specified offset.
    ///
    /// The offset is assumed to be an offset into the view of the matrix where extension field
    /// elements are decomposed into base field elements. This offset must be compatible with the
    /// values supplied into [Matrix::get_base_element()] method.
    ///
    /// Evaluation is performed over the domain specified by the provided twiddles and offsets.
    ///
    /// # Panics
    /// Panics if:
    /// - `poly_offset` greater than or equal to the number of base field columns in `polys`.
    /// - Number of offsets is not a power of two.
    /// - Number of offsets is smaller than or equal to the polynomial size.
    /// - The number of twiddles is not half the size of the polynomial size.
    /// - Number of offsets is smaller than the length of the data buffer
    pub fn new_with_buffer<E>(
        data_buffer: Vec<[B; N]>,
        polys: &ColMatrix<E>,
        poly_offset: usize,
        offsets: &[B],
        twiddles: &[B],
    ) -> Self
    where
        E: FieldElement<BaseField = B>,
    {
        let poly_size = polys.num_rows();
        let domain_size = offsets.len();
        let mut data = data_buffer;

        assert!(domain_size.is_power_of_two());
        assert!(domain_size > poly_size);
        assert_eq!(poly_size, twiddles.len() * 2);
        assert!(poly_offset < polys.num_base_cols());
        assert_eq!(data.len(), domain_size);

        // determine the number of polynomials to add to this segment; this number can be either N,
        // or smaller than N when there are fewer than N polynomials remaining to be processed
        let num_polys_remaining = polys.num_base_cols() - poly_offset;
        let num_polys = if num_polys_remaining < N {
            num_polys_remaining
        } else {
            N
        };

        // evaluate the polynomials either in a single thread or multiple threads, depending
        // on whether `concurrent` feature is enabled and domain size is greater than 1024;

        if cfg!(feature = "concurrent") && domain_size >= MIN_CONCURRENT_SIZE {
            #[cfg(feature = "concurrent")]
            data.par_chunks_mut(poly_size)
                .zip(offsets.par_chunks(poly_size))
                .for_each(|(d_chunk, o_chunk)| {
                    // TODO: investigate multi-threaded copy
                    if num_polys == N {
                        Self::copy_polys(d_chunk, polys, poly_offset, o_chunk);
                    } else {
                        Self::copy_polys_partial(d_chunk, polys, poly_offset, num_polys, o_chunk);
                    }
                    concurrent::split_radix_fft(d_chunk, twiddles);
                });
            #[cfg(feature = "concurrent")]
            concurrent::permute(&mut data);
        } else {
            data.chunks_mut(poly_size)
                .zip(offsets.chunks(poly_size))
                .for_each(|(d_chunk, o_chunk)| {
                    if num_polys == N {
                        Self::copy_polys(d_chunk, polys, poly_offset, o_chunk);
                    } else {
                        Self::copy_polys_partial(d_chunk, polys, poly_offset, num_polys, o_chunk);
                    }
                    d_chunk.fft_in_place(twiddles);
                });
            data.permute();
        }

        Segment { data }
    }

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

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

    /// Returns the underlying vector of arrays for this segment.
    pub fn into_data(self) -> Vec<[B; N]> {
        self.data
    }

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

    /// Copies N polynomials starting at the specified base column offset (`poly_offset`) into the
    /// specified destination. Each polynomial coefficient is offset by the specified offset.
    fn copy_polys<E: FieldElement<BaseField = B>>(
        dest: &mut [[B; N]],
        polys: &ColMatrix<E>,
        poly_offset: usize,
        offsets: &[B],
    ) {
        for row_idx in 0..dest.len() {
            for i in 0..N {
                let coeff = polys.get_base_element(poly_offset + i, row_idx);
                dest[row_idx][i] = coeff * offsets[row_idx];
            }
        }
    }

    /// Similar to `clone_and_shift` method above, but copies `num_polys` polynomials instead of
    /// `N` polynomials.
    ///
    /// Assumes that `num_polys` is smaller than `N`.
    fn copy_polys_partial<E: FieldElement<BaseField = B>>(
        dest: &mut [[B; N]],
        polys: &ColMatrix<E>,
        poly_offset: usize,
        num_polys: usize,
        offsets: &[B],
    ) {
        debug_assert!(num_polys < N);
        for row_idx in 0..dest.len() {
            for i in 0..num_polys {
                let coeff = polys.get_base_element(poly_offset + i, row_idx);
                dest[row_idx][i] = coeff * offsets[row_idx];
            }
        }
    }
}

impl<B: StarkField, const N: usize> Deref for Segment<B, N> {
    type Target = Vec<[B; N]>;

    fn deref(&self) -> &Self::Target {
        &self.data
    }
}

// CONCURRENT FFT IMPLEMENTATION
// ================================================================================================

/// Multi-threaded implementations of FFT and permutation algorithms. These are very similar to
/// the functions implemented in `winter-math::fft::concurrent` module, but are adapted to work
/// with slices of element arrays.
#[cfg(feature = "concurrent")]
mod concurrent {
    use super::{FftInputs, StarkField};
    use math::fft::permute_index;
    use utils::{iterators::*, rayon};

    /// In-place recursive FFT with permuted output.
    /// Adapted from: https://github.com/0xProject/OpenZKP/tree/master/algebra/primefield/src/fft
    pub fn split_radix_fft<B: StarkField, const N: usize>(data: &mut [[B; N]], twiddles: &[B]) {
        // generator of the domain should be in the middle of twiddles
        let n = data.len();
        let g = twiddles[twiddles.len() / 2];
        debug_assert_eq!(g.exp((n as u32).into()), B::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(data, inner_len, stretch);

        // apply inner FFTs
        data.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(data, inner_len, stretch);

        // apply outer FFTs
        data.par_chunks_mut(outer_len)
            .enumerate()
            .for_each(|(i, row)| {
                if i > 0 {
                    let i = permute_index(inner_len, i);
                    let inner_twiddle = g.exp_vartime((i as u32).into());
                    let mut outer_twiddle = inner_twiddle;
                    for element in row.iter_mut().skip(1) {
                        for col_idx in 0..N {
                            element[col_idx] = element[col_idx] * outer_twiddle;
                        }
                        outer_twiddle = outer_twiddle * inner_twiddle;
                    }
                }
                row.fft_in_place(&twiddles)
            });
    }

    // PERMUTATIONS
    // --------------------------------------------------------------------------------------------

    pub fn permute<T: Send>(v: &mut [T]) {
        let n = v.len();
        let num_batches = rayon::current_num_threads().next_power_of_two() * 2;
        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 [T]) };
                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 = permute_index(n, i);
                        if j > i {
                            values.swap(i, j);
                        }
                    }
                });
            }
        });
    }

    // TRANSPOSING
    // --------------------------------------------------------------------------------------------

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

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

        // iterate over upper-left triangle, working in 2x2 blocks
        // TODO: investigate concurrent implementation
        for row in (0..size).step_by(2) {
            let i = row * size + row;
            data.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;
                data.swap(i, j);
                data.swap(i + 1, j + size);
                data.swap(i + size, j + 1);
                data.swap(i + size + 1, j + size + 1);
            }
        }
    }

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

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


// 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::{
    math::{fields::f64::BaseElement, get_power_series, polynom, StarkField},
    ColMatrix, RowMatrix,
};
use rand_utils::rand_vector;
use utils::collections::Vec;

#[test]
fn test_eval_poly_with_offset_matrix() {
    let n = 256;
    let num_polys = 64;
    let blowup_factor = 8;

    // generate random columns. Each column is a polynomial of degree n - 1.
    let mut columns: Vec<Vec<BaseElement>> = (0..num_polys).map(|_| rand_vector(n)).collect();

    // evaluate columns using the row matrix implementation.
    let row_matrix =
        RowMatrix::evaluate_polys::<8>(&ColMatrix::new(columns.clone()), blowup_factor);

    // evaluate columns using the using the polynomial evaluation implementation.
    let offset = BaseElement::GENERATOR;
    let domain = build_domain(n * blowup_factor);
    let shifted_domain = domain.iter().map(|&x| x * offset).collect::<Vec<_>>();
    for p in columns.iter_mut() {
        *p = polynom::eval_many(p, &shifted_domain);
    }

    // compare the results of the two implementations row by row.
    for row in 0..n * blowup_factor {
        let row_matrix_row = row_matrix.row(row);
        let eval_col_row = get_row(&columns, row);
        assert_eq!(row_matrix_row, eval_col_row);
    }
}

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

/// Builds a domain of size `size` using the primitive element of the field.
fn build_domain(size: usize) -> Vec<BaseElement> {
    let g = BaseElement::get_root_of_unity(size.ilog2());
    get_power_series(g, size)
}

/// Returns a row of the column major matrix.
fn get_row(columns: &[Vec<BaseElement>], row_id: usize) -> Vec<BaseElement> {
    columns.iter().map(|col| col[row_id]).collect()
}


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

use crate::TraceTable;
use air::{
    Air, AirContext, Assertion, EvaluationFrame, FieldExtension, ProofOptions, TraceInfo,
    TransitionConstraintDegree,
};
use math::{fields::f128::BaseElement, FieldElement, StarkField};
use utils::collections::Vec;

// FIBONACCI TRACE BUILDER
// ================================================================================================

pub fn build_fib_trace(length: usize) -> TraceTable<BaseElement> {
    assert!(length.is_power_of_two(), "length must be a power of 2");

    let mut reg1 = vec![BaseElement::ONE];
    let mut reg2 = vec![BaseElement::ONE];

    for i in 0..(length / 2 - 1) {
        reg1.push(reg1[i] + reg2[i]);
        reg2.push(reg1[i] + BaseElement::from(2u8) * reg2[i]);
    }

    TraceTable::init(vec![reg1, reg2])
}

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

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

impl MockAir {
    pub fn with_trace_length(trace_length: usize) -> Self {
        Self::new(
            TraceInfo::new(4, trace_length),
            (),
            ProofOptions::new(32, 8, 0, FieldExtension::None, 4, 31),
        )
    }

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

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

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

    fn new(trace_info: TraceInfo, _pub_inputs: (), _options: ProofOptions) -> Self {
        let context = build_context(trace_info, 8, 1);
        MockAir {
            context,
            assertions: Vec::new(),
            periodic_columns: Vec::new(),
        }
    }

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

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

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

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

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

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


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

use super::TraceLde;

// TRACE COMMITMENT
// ================================================================================================

/// Execution trace commitment.
///
/// The describes one or more trace segments, each consisting of the following components:
/// * Evaluations of a trace segment's polynomials over the LDE domain.
/// * Merkle tree where each leaf in the tree corresponds to a row in the trace LDE matrix.
pub struct TraceCommitment<E: FieldElement, H: ElementHasher<BaseField = E::BaseField>> {
    trace_lde: TraceLde<E>,
    main_segment_tree: MerkleTree<H>,
    aux_segment_trees: Vec<MerkleTree<H>>,
}

impl<E: FieldElement, H: ElementHasher<BaseField = E::BaseField>> TraceCommitment<E, H> {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Creates a new trace commitment from the provided main trace low-degree extension and the
    /// corresponding Merkle tree commitment.
    pub fn new(
        main_trace_lde: RowMatrix<E::BaseField>,
        main_trace_tree: MerkleTree<H>,
        blowup: usize,
    ) -> Self {
        assert_eq!(
            main_trace_lde.num_rows(),
            main_trace_tree.leaves().len(),
            "number of rows in trace LDE must be the same as number of leaves in trace commitment"
        );
        Self {
            trace_lde: TraceLde::new(main_trace_lde, blowup),
            main_segment_tree: main_trace_tree,
            aux_segment_trees: Vec::new(),
        }
    }

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

    /// Adds the provided auxiliary segment trace LDE and Merkle tree to this trace commitment.
    pub fn add_segment(&mut self, aux_segment_lde: RowMatrix<E>, aux_segment_tree: MerkleTree<H>) {
        assert_eq!(
            aux_segment_lde.num_rows(),
            aux_segment_tree.leaves().len(),
            "number of rows in trace LDE must be the same as number of leaves in trace commitment"
        );

        self.trace_lde.add_aux_segment(aux_segment_lde);
        self.aux_segment_trees.push(aux_segment_tree);
    }

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

    /// Returns the execution trace for this commitment.
    ///
    /// The trace contains both the main trace segment and the auxiliary trace segments (if any).
    pub fn trace_table(&self) -> &TraceLde<E> {
        &self.trace_lde
    }

    // QUERY TRACE
    // --------------------------------------------------------------------------------------------
    /// Returns trace table rows at the specified positions along with Merkle authentication paths
    /// from the commitment root to these rows.
    pub fn query(&self, positions: &[usize]) -> Vec<Queries> {
        // build queries for the main trace segment
        let mut result = vec![build_segment_queries(
            self.trace_lde.get_main_segment(),
            &self.main_segment_tree,
            positions,
        )];

        // build queries for auxiliary trace segments
        for (i, segment_tree) in self.aux_segment_trees.iter().enumerate() {
            let segment_lde = self.trace_lde.get_aux_segment(i);
            result.push(build_segment_queries(segment_lde, segment_tree, positions));
        }

        result
    }

    // TEST HELPERS
    // --------------------------------------------------------------------------------------------

    /// Returns the root of the commitment Merkle tree.
    #[cfg(test)]
    pub fn main_trace_root(&self) -> H::Digest {
        *self.main_segment_tree.root()
    }

    /// Returns the entire trace for the column at the specified index.
    #[cfg(test)]
    pub fn get_main_trace_column(&self, col_idx: usize) -> Vec<E::BaseField> {
        let trace = self.trace_lde.get_main_segment();
        (0..trace.num_rows())
            .map(|row_idx| trace.get(col_idx, row_idx))
            .collect()
    }
}

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

fn build_segment_queries<E, H>(
    segment_lde: &RowMatrix<E>,
    segment_tree: &MerkleTree<H>,
    positions: &[usize],
) -> Queries
where
    E: FieldElement,
    H: ElementHasher<BaseField = E::BaseField>,
{
    // for each position, get the corresponding row from the trace segment LDE and put all these
    // rows into a single vector
    let trace_states = positions
        .iter()
        .map(|&pos| segment_lde.row(pos).to_vec())
        .collect::<Vec<_>>();

    // build Merkle authentication paths to the leaves specified by positions
    let trace_proof = segment_tree
        .prove_batch(positions)
        .expect("failed to generate a Merkle proof for trace queries");

    Queries::new(trace_proof, trace_states)
}


// 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::{matrix::MultiColumnIter, ColMatrix};
use air::{Air, AuxTraceRandElements, EvaluationFrame, TraceInfo, TraceLayout};
use math::{polynom, FieldElement, StarkField};

mod trace_lde;
pub use trace_lde::TraceLde;

mod poly_table;
pub use poly_table::TracePolyTable;

mod trace_table;
pub use trace_table::{TraceTable, TraceTableFragment};

mod commitment;
pub use commitment::TraceCommitment;

#[cfg(test)]
mod tests;

// TRACE TRAIT
// ================================================================================================
/// Defines an execution trace of a computation.
///
/// Execution trace can be reduced to a two-dimensional matrix in which each row represents the
/// state of a computation at a single point in time and each column corresponds to an algebraic
/// column tracked over all steps of the computation.
///
/// Building a trace is required for STARK proof generation. An execution trace of a specific
/// instance of a computation must be supplied to [Prover::prove()](super::Prover::prove) method
/// to generate a STARK proof.
///
/// This crate exposes one concrete implementation of the [Trace] trait: [TraceTable]. This
/// implementation supports concurrent trace generation and should be sufficient in most
/// situations. However, if functionality provided by [TraceTable] is not sufficient, uses can
/// provide custom implementations of the [Trace] trait which better suit their needs.
pub trait Trace: Sized {
    /// Base field for this execution trace.
    ///
    /// All cells of this execution trace contain values which are elements in this field.
    type BaseField: StarkField;

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

    /// Returns a description of how columns of this trace are arranged into trace segments.
    fn layout(&self) -> &TraceLayout;

    /// Returns the number of rows in this trace.
    fn length(&self) -> usize;

    /// Returns metadata associated with this trace.
    fn meta(&self) -> &[u8];

    /// Returns a reference to a [Matrix] describing the main segment of this trace.
    fn main_segment(&self) -> &ColMatrix<Self::BaseField>;

    /// Builds and returns the next auxiliary trace segment. If there are no more segments to
    /// build (i.e., the trace is complete), None is returned.
    ///
    /// The `aux_segments` slice contains a list of auxiliary trace segments built as a result
    /// of prior invocations of this function. Thus, for example, on the first invocation,
    /// `aux_segments` will be empty; on the second invocation, it will contain a single matrix
    /// (the one built during the first invocation) etc.
    fn build_aux_segment<E: FieldElement<BaseField = Self::BaseField>>(
        &mut self,
        aux_segments: &[ColMatrix<E>],
        rand_elements: &[E],
    ) -> Option<ColMatrix<E>>;

    /// Reads an evaluation frame from the main trace segment at the specified row.
    fn read_main_frame(&self, row_idx: usize, frame: &mut EvaluationFrame<Self::BaseField>);

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

    /// Returns trace info for this trace.
    fn get_info(&self) -> TraceInfo {
        TraceInfo::new_multi_segment(self.layout().clone(), self.length(), self.meta().to_vec())
    }

    /// Returns the number of columns in the main segment of this trace.
    fn main_trace_width(&self) -> usize {
        self.layout().main_trace_width()
    }

    /// Returns the number of columns in all auxiliary trace segments.
    fn aux_trace_width(&self) -> usize {
        self.layout().aux_trace_width()
    }

    // VALIDATION
    // --------------------------------------------------------------------------------------------
    /// Checks if this trace is valid against the specified AIR, and panics if not.
    ///
    /// NOTE: this is a very expensive operation and is intended for use only in debug mode.
    fn validate<A, E>(
        &self,
        air: &A,
        aux_segments: &[ColMatrix<E>],
        aux_rand_elements: &AuxTraceRandElements<E>,
    ) where
        A: Air<BaseField = Self::BaseField>,
        E: FieldElement<BaseField = Self::BaseField>,
    {
        // make sure the width align; if they don't something went terribly wrong
        assert_eq!(
            self.main_trace_width(),
            air.trace_layout().main_trace_width(),
            "inconsistent trace width: expected {}, but was {}",
            self.main_trace_width(),
            air.trace_layout().main_trace_width(),
        );

        // --- 1. make sure the assertions are valid ----------------------------------------------

        // first, check assertions against the main segment of the execution trace
        for assertion in air.get_assertions() {
            assertion.apply(self.length(), |step, value| {
                assert!(
                    value == self.main_segment().get(assertion.column(), step),
                    "trace does not satisfy assertion main_trace({}, {}) == {}",
                    assertion.column(),
                    step,
                    value
                );
            });
        }

        // then, check assertions against auxiliary trace segments
        for assertion in air.get_aux_assertions(aux_rand_elements) {
            // find which segment the assertion is for and remap assertion column index to the
            // column index in the context of this segment
            let mut column_idx = assertion.column();
            let mut segment_idx = 0;
            for i in 0..self.layout().num_aux_segments() {
                let segment_width = self.layout().get_aux_segment_width(i);
                if column_idx < segment_width {
                    segment_idx = i;
                    break;
                }
                column_idx -= segment_width;
            }

            // get the matrix and verify the assertion against it
            assertion.apply(self.length(), |step, value| {
                assert!(
                    value == aux_segments[segment_idx].get(column_idx, step),
                    "trace does not satisfy assertion aux_trace({}, {}) == {}",
                    assertion.column(),
                    step,
                    value
                );
            });
        }

        // --- 2. make sure this trace satisfies all transition constraints -----------------------

        // collect the info needed to build periodic values for a specific step
        let g = air.trace_domain_generator();
        let periodic_values_polys = air.get_periodic_column_polys();
        let mut periodic_values = vec![Self::BaseField::ZERO; periodic_values_polys.len()];

        // initialize buffers to hold evaluation frames and results of constraint evaluations
        let mut x = Self::BaseField::ONE;
        let mut main_frame = EvaluationFrame::new(self.main_trace_width());
        let mut aux_frame = if air.trace_info().is_multi_segment() {
            Some(EvaluationFrame::<E>::new(self.aux_trace_width()))
        } else {
            None
        };
        let mut main_evaluations =
            vec![Self::BaseField::ZERO; air.context().num_main_transition_constraints()];
        let mut aux_evaluations = vec![E::ZERO; air.context().num_aux_transition_constraints()];

        // we check transition constraints on all steps except the last k steps, where k is the
        // number of steps exempt from transition constraints (guaranteed to be at least 1)
        for step in 0..self.length() - air.context().num_transition_exemptions() {
            // build periodic values
            for (p, v) in periodic_values_polys.iter().zip(periodic_values.iter_mut()) {
                let num_cycles = air.trace_length() / p.len();
                let x = x.exp((num_cycles as u32).into());
                *v = polynom::eval(p, x);
            }

            // evaluate transition constraints for the main trace segment and make sure they all
            // evaluate to zeros
            self.read_main_frame(step, &mut main_frame);
            air.evaluate_transition(&main_frame, &periodic_values, &mut main_evaluations);
            for (i, &evaluation) in main_evaluations.iter().enumerate() {
                assert!(
                    evaluation == Self::BaseField::ZERO,
                    "main transition constraint {i} did not evaluate to ZERO at step {step}"
                );
            }

            // evaluate transition constraints for auxiliary trace segments (if any) and make
            // sure they all evaluate to zeros
            if let Some(ref mut aux_frame) = aux_frame {
                read_aux_frame(aux_segments, step, aux_frame);
                air.evaluate_aux_transition(
                    &main_frame,
                    aux_frame,
                    &periodic_values,
                    aux_rand_elements,
                    &mut aux_evaluations,
                );
                for (i, &evaluation) in aux_evaluations.iter().enumerate() {
                    assert!(
                        evaluation == E::ZERO,
                        "auxiliary transition constraint {i} did not evaluate to ZERO at step {step}"
                    );
                }
            }

            // update x coordinate of the domain
            x *= g;
        }
    }
}

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

/// Reads an evaluation frame from the set of provided auxiliary segments. This expects that
/// `aux_segments` contains at least one entry.
///
/// This is probably not the most efficient implementation, but since we call this function only
/// for trace validation purposes (which is done in debug mode only), we don't care all that much
/// about its performance.
fn read_aux_frame<E>(aux_segments: &[ColMatrix<E>], row_idx: usize, frame: &mut EvaluationFrame<E>)
where
    E: FieldElement,
{
    for (column, current_value) in MultiColumnIter::new(aux_segments).zip(frame.current_mut()) {
        *current_value = column[row_idx];
    }

    let next_row_idx = (row_idx + 1) % aux_segments[0].num_rows();
    for (column, next_value) in MultiColumnIter::new(aux_segments).zip(frame.next_mut()) {
        *next_value = column[next_row_idx];
    }
}


// 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::{
    matrix::{ColumnIter, MultiColumnIter},
    ColMatrix,
};
use math::{FieldElement, StarkField};
use utils::collections::Vec;

// TRACE POLYNOMIAL TABLE
// ================================================================================================

/// Trace polynomials in coefficient from for all segments of the execution trace.
///
/// Coefficients of the polynomials for the main trace segment are always in the base field.
/// However, coefficients of the polynomials for the auxiliary trace segments may be either in the
/// base field, or in the extension field, depending on whether extension field is being used.
pub struct TracePolyTable<E: FieldElement> {
    main_segment_polys: ColMatrix<E::BaseField>,
    aux_segment_polys: Vec<ColMatrix<E>>,
}

impl<E: FieldElement> TracePolyTable<E> {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Creates a new table of trace polynomials from the provided main trace segment polynomials.
    pub fn new(main_trace_polys: ColMatrix<E::BaseField>) -> Self {
        Self {
            main_segment_polys: main_trace_polys,
            aux_segment_polys: Vec::new(),
        }
    }

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

    /// Adds the provided auxiliary segment polynomials to this polynomial table.
    pub fn add_aux_segment(&mut self, aux_segment_polys: ColMatrix<E>) {
        assert_eq!(
            self.main_segment_polys.num_rows(),
            aux_segment_polys.num_rows(),
            "polynomials in auxiliary segment must be of the same size as in the main segment"
        );
        self.aux_segment_polys.push(aux_segment_polys);
    }

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

    /// Returns the size of each polynomial - i.e. size of a vector needed to hold a polynomial.
    pub fn poly_size(&self) -> usize {
        self.main_segment_polys.num_rows()
    }

    /// Evaluates all trace polynomials (across all trace segments) at the specified point `x`.
    pub fn evaluate_at(&self, x: E) -> Vec<E> {
        let mut result = self.main_segment_polys.evaluate_columns_at(x);
        for aux_polys in self.aux_segment_polys.iter() {
            result.append(&mut aux_polys.evaluate_columns_at(x));
        }
        result
    }

    /// Returns an out-of-domain evaluation frame constructed by evaluating trace polynomials
    /// for all columns at points z and z * g, where g is the generator of the trace domain.
    pub fn get_ood_frame(&self, z: E) -> Vec<Vec<E>> {
        let g = E::from(E::BaseField::get_root_of_unity(self.poly_size().ilog2()));
        vec![self.evaluate_at(z), self.evaluate_at(z * g)]
    }

    /// Returns an iterator over the polynomials of the main trace segment.
    pub fn main_trace_polys(&self) -> ColumnIter<E::BaseField> {
        self.main_segment_polys.columns()
    }

    /// Returns an iterator over the polynomials of all auxiliary trace segments.
    pub fn aux_trace_polys(&self) -> MultiColumnIter<E> {
        MultiColumnIter::new(self.aux_segment_polys.as_slice())
    }

    // TEST HELPERS
    // --------------------------------------------------------------------------------------------

    /// Returns the number of polynomials in the main segment of the trace.
    #[cfg(test)]
    pub fn num_main_trace_polys(&self) -> usize {
        self.main_segment_polys.num_cols()
    }

    /// Returns a polynomial from the main segment of the trace at the specified index.
    #[cfg(test)]
    pub fn get_main_trace_poly(&self, idx: usize) -> &[E::BaseField] {
        self.main_segment_polys.get_column(idx)
    }
}


// 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::{
    tests::{build_fib_trace, MockAir},
    trace::TracePolyTable,
    RowMatrix, StarkDomain, Trace, TraceCommitment,
};
use crypto::{hashers::Blake3_256, ElementHasher, MerkleTree};
use math::{
    fields::f128::BaseElement, get_power_series, get_power_series_with_offset, polynom,
    FieldElement, StarkField,
};
use utils::collections::Vec;

type Blake3 = Blake3_256<BaseElement>;

#[test]
fn new_trace_table() {
    let trace_length = 8;
    let trace = build_fib_trace(trace_length * 2);

    assert_eq!(2, trace.main_trace_width());
    assert_eq!(8, trace.length());

    let expected: Vec<BaseElement> = vec![1u32, 2, 5, 13, 34, 89, 233, 610]
        .into_iter()
        .map(BaseElement::from)
        .collect();
    assert_eq!(expected, trace.get_column(0));

    let expected: Vec<BaseElement> = vec![1u32, 3, 8, 21, 55, 144, 377, 987]
        .into_iter()
        .map(BaseElement::from)
        .collect();
    assert_eq!(expected, trace.get_column(1));
}

#[test]
fn extend_trace_table() {
    // build the trace and the domain
    let trace_length = 8;
    let air = MockAir::with_trace_length(trace_length);
    let trace = build_fib_trace(trace_length * 2);
    let domain = StarkDomain::new(&air);

    // build extended trace commitment
    let trace_polys = trace.main_segment().interpolate_columns();
    let trace_lde = RowMatrix::evaluate_polys_over::<8>(&trace_polys, &domain);
    let trace_tree = trace_lde.commit_to_rows::<Blake3>();
    let trace_comm = TraceCommitment::<BaseElement, Blake3>::new(
        trace_lde,
        trace_tree,
        domain.trace_to_lde_blowup(),
    );
    let trace_polys = TracePolyTable::<BaseElement>::new(trace_polys);

    assert_eq!(2, trace_comm.trace_table().main_trace_width());
    assert_eq!(64, trace_comm.trace_table().trace_len());

    // make sure trace polynomials evaluate to Fibonacci trace
    let trace_root = BaseElement::get_root_of_unity(trace_length.ilog2());
    let trace_domain = get_power_series(trace_root, trace_length);
    assert_eq!(2, trace_polys.num_main_trace_polys());
    assert_eq!(
        vec![1u32, 2, 5, 13, 34, 89, 233, 610]
            .into_iter()
            .map(BaseElement::from)
            .collect::<Vec<BaseElement>>(),
        polynom::eval_many(trace_polys.get_main_trace_poly(0), &trace_domain)
    );
    assert_eq!(
        vec![1u32, 3, 8, 21, 55, 144, 377, 987]
            .into_iter()
            .map(BaseElement::from)
            .collect::<Vec<BaseElement>>(),
        polynom::eval_many(trace_polys.get_main_trace_poly(1), &trace_domain)
    );

    // make sure column values are consistent with trace polynomials
    let lde_domain = build_lde_domain(domain.lde_domain_size());
    assert_eq!(
        trace_polys.get_main_trace_poly(0),
        polynom::interpolate(&lde_domain, &trace_comm.get_main_trace_column(0), true)
    );
    assert_eq!(
        trace_polys.get_main_trace_poly(1),
        polynom::interpolate(&lde_domain, &trace_comm.get_main_trace_column(1), true)
    );
}

#[test]
fn commit_trace_table() {
    // build the trade and the domain
    let trace_length = 8;
    let air = MockAir::with_trace_length(trace_length);
    let trace = build_fib_trace(trace_length * 2);
    let domain = StarkDomain::new(&air);

    // build extended trace commitment
    let trace_polys = trace.main_segment().interpolate_columns();
    let trace_lde = RowMatrix::evaluate_polys_over::<8>(&trace_polys, &domain);
    let trace_tree = trace_lde.commit_to_rows::<Blake3>();
    let trace_comm = TraceCommitment::<BaseElement, Blake3>::new(
        trace_lde,
        trace_tree,
        domain.trace_to_lde_blowup(),
    );

    // build Merkle tree from trace rows
    let trace_table = trace_comm.trace_table();
    let mut hashed_states = Vec::new();
    let mut trace_state = vec![BaseElement::ZERO; trace_table.main_trace_width()];
    #[allow(clippy::needless_range_loop)]
    for i in 0..trace_table.trace_len() {
        for j in 0..trace_table.main_trace_width() {
            trace_state[j] = trace_table.get_main_segment().get(j, i);
        }
        let buf = Blake3::hash_elements(&trace_state);
        hashed_states.push(buf);
    }
    let expected_tree = MerkleTree::<Blake3>::new(hashed_states).unwrap();

    // compare the result
    assert_eq!(*expected_tree.root(), trace_comm.main_trace_root())
}

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

fn build_lde_domain<B: StarkField>(domain_size: usize) -> Vec<B> {
    let g = B::get_root_of_unity(domain_size.ilog2());
    get_power_series_with_offset(g, B::GENERATOR, domain_size)
}


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

use crate::RowMatrix;
use air::EvaluationFrame;
use math::FieldElement;
use utils::collections::Vec;

// TRACE LOW DEGREE EXTENSION
// ================================================================================================
/// Contains all segments of the extended execution trace.
///
/// Segments are stored in two groups:
/// - Main segment: this is the first trace segment generated by the prover. Values in this segment
///   will always be elements in the base field (even when an extension field is used).
/// - Auxiliary segments: a list of 0 or more segments for traces generated after the prover
///   commits to the first trace segment. Currently, at most 1 auxiliary segment is possible.
pub struct TraceLde<E: FieldElement> {
    main_segment_lde: RowMatrix<E::BaseField>,
    aux_segment_ldes: Vec<RowMatrix<E>>,
    blowup: usize,
}

impl<E: FieldElement> TraceLde<E> {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Creates a new trace low-degree extension table from the provided main trace segment LDE.
    pub fn new(main_trace_lde: RowMatrix<E::BaseField>, blowup: usize) -> Self {
        Self {
            main_segment_lde: main_trace_lde,
            aux_segment_ldes: Vec::new(),
            blowup,
        }
    }

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

    /// Adds the provided auxiliary segment LDE to this trace LDE.
    pub fn add_aux_segment(&mut self, aux_segment_lde: RowMatrix<E>) {
        assert_eq!(
            self.main_segment_lde.num_rows(),
            aux_segment_lde.num_rows(),
            "number of rows in auxiliary segment must be of the same as in the main segment"
        );
        self.aux_segment_ldes.push(aux_segment_lde);
    }

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

    /// Returns number of columns in the main segment of the execution trace.
    pub fn main_trace_width(&self) -> usize {
        self.main_segment_lde.num_cols()
    }

    /// Returns number of columns in the auxiliary segments of the execution trace.
    pub fn aux_trace_width(&self) -> usize {
        self.aux_segment_ldes
            .iter()
            .fold(0, |s, m| s + m.num_cols())
    }

    /// Returns the number of rows in the execution trace.
    pub fn trace_len(&self) -> usize {
        self.main_segment_lde.num_rows()
    }

    /// Returns blowup factor which was used to extend original execution trace into trace LDE.
    pub fn blowup(&self) -> usize {
        self.blowup
    }

    /// Reads current and next rows from the main trace segment into the specified frame.
    pub fn read_main_trace_frame_into(
        &self,
        lde_step: usize,
        frame: &mut EvaluationFrame<E::BaseField>,
    ) {
        // at the end of the trace, next state wraps around and we read the first step again
        let next_lde_step = (lde_step + self.blowup()) % self.trace_len();

        // copy main trace segment values into the frame
        frame
            .current_mut()
            .copy_from_slice(self.main_segment_lde.row(lde_step));
        frame
            .next_mut()
            .copy_from_slice(self.main_segment_lde.row(next_lde_step));
    }

    /// Reads current and next rows from the auxiliary trace segment into the specified frame.
    ///
    /// # Panics
    /// This currently assumes that there is exactly one auxiliary trace segment, and will panic
    /// otherwise.
    pub fn read_aux_trace_frame_into(&self, lde_step: usize, frame: &mut EvaluationFrame<E>) {
        // at the end of the trace, next state wraps around and we read the first step again
        let next_lde_step = (lde_step + self.blowup()) % self.trace_len();

        // copy auxiliary trace segment values into the frame
        let segment = &self.aux_segment_ldes[0];
        frame.current_mut().copy_from_slice(segment.row(lde_step));
        frame.next_mut().copy_from_slice(segment.row(next_lde_step));
    }

    /// Returns a reference to [Matrix] representing the main trace segment.
    pub fn get_main_segment(&self) -> &RowMatrix<E::BaseField> {
        &self.main_segment_lde
    }

    /// Returns a reference to a [Matrix] representing an auxiliary trace segment at the specified
    /// index.
    pub fn get_aux_segment(&self, aux_segment_idx: usize) -> &RowMatrix<E> {
        &self.aux_segment_ldes[aux_segment_idx]
    }
}


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

use super::{ColMatrix, Trace};
use air::{EvaluationFrame, TraceInfo, TraceLayout};
use math::{FieldElement, StarkField};
use utils::{collections::Vec, uninit_vector};

#[cfg(not(feature = "concurrent"))]
use utils::collections::vec;

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

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

const MIN_FRAGMENT_LENGTH: usize = 2;

// TRACE TABLE
// ================================================================================================
/// A concrete implementation of the [Trace] trait.
///
/// This implementation supports concurrent trace generation and should be sufficient for most use
/// cases. There are two ways to create a trace table trace.
///
/// First, you can use the [TraceTable::init()] function which takes a set of vectors as a
/// parameter, where each vector contains values for a given column of the trace. This approach
/// allows you to build an execution trace as you see fit, as long as it meets a basic set of
/// requirements. These requirements are:
///
/// 1. Lengths of all columns in the execution trace must be the same.
/// 2. The length of the columns must be some power of two.
///
/// The other approach is to use [TraceTable::new()] function, which takes trace width and
/// length as parameters. This function will allocate memory for the trace, but will not fill it
/// with data. To fill the execution trace, you can use the [fill()](TraceTable::fill) method,
/// which takes two closures as parameters:
///
/// 1. The first closure is responsible for initializing the first state of the computation
///    (the first row of the execution trace).
/// 2. The second closure receives the previous state of the execution trace as input, and must
///    update it to the next state of the computation.
///
/// You can also use [TraceTable::with_meta()] function to create a blank execution trace.
/// This function work just like [TraceTable::new()] function, but also takes a metadata
/// parameter which can be an arbitrary sequence of bytes up to 64KB in size.
///
/// # Concurrent trace generation
/// For computations which consist of many small independent computations, we can generate the
/// execution trace of the entire computation by building fragments of the trace in parallel,
/// and then joining these fragments together.
///
/// For this purpose, `TraceTable` struct exposes [fragments()](TraceTable::fragments)
/// method, which takes fragment length as a parameter, breaks the execution trace into equally
/// sized fragments, and returns an iterator over these fragments. You can then use fragment's
/// [fill()](TraceTableFragment::fill) method to fill all fragments with data in parallel.
/// The semantics of the fragment's [TraceTableFragment::fill()] method are identical to the
/// semantics of the [TraceTable::fill()] method.
pub struct TraceTable<B: StarkField> {
    layout: TraceLayout,
    trace: ColMatrix<B>,
    meta: Vec<u8>,
}

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

    /// Creates a new execution trace of the specified width and length.
    ///
    /// This allocates all the required memory for the trace, but does not initialize it. It is
    /// expected that the trace will be filled using one of the data mutator methods.
    ///
    /// # Panics
    /// Panics if:
    /// * `width` is zero or greater than 255.
    /// * `length` is smaller than 8, greater than biggest multiplicative subgroup in the field
    ///   `B`, or is not a power of two.
    pub fn new(width: usize, length: usize) -> Self {
        Self::with_meta(width, length, vec![])
    }

    /// Creates a new execution trace of the specified width and length, and with the specified
    /// metadata.
    ///
    /// This allocates all the required memory for the trace, but does not initialize it. It is
    /// expected that the trace will be filled using one of the data mutator methods.
    ///
    /// # Panics
    /// Panics if:
    /// * `width` is zero or greater than 255.
    /// * `length` is smaller than 8, greater than the biggest multiplicative subgroup in the
    ///   field `B`, or is not a power of two.
    /// * Length of `meta` is greater than 65535;
    pub fn with_meta(width: usize, length: usize, meta: Vec<u8>) -> Self {
        assert!(
            width > 0,
            "execution trace must consist of at least one column"
        );
        assert!(
            width <= TraceInfo::MAX_TRACE_WIDTH,
            "execution trace width cannot be greater than {}, but was {}",
            TraceInfo::MAX_TRACE_WIDTH,
            width
        );
        assert!(
            length >= TraceInfo::MIN_TRACE_LENGTH,
            "execution trace must be at least {} steps long, but was {}",
            TraceInfo::MIN_TRACE_LENGTH,
            length
        );
        assert!(
            length.is_power_of_two(),
            "execution trace length must be a power of 2"
        );
        assert!(
            length.ilog2() <= B::TWO_ADICITY,
            "execution trace length cannot exceed 2^{} steps, but was 2^{}",
            B::TWO_ADICITY,
            length.ilog2()
        );
        assert!(
            meta.len() <= TraceInfo::MAX_META_LENGTH,
            "number of metadata bytes cannot be greater than {}, but was {}",
            TraceInfo::MAX_META_LENGTH,
            meta.len()
        );

        let columns = unsafe { (0..width).map(|_| uninit_vector(length)).collect() };
        Self {
            layout: TraceLayout::new(width, [0], [0]),
            trace: ColMatrix::new(columns),
            meta,
        }
    }

    /// Creates a new execution trace from a list of provided trace columns.
    ///
    /// # Panics
    /// Panics if:
    /// * The `columns` vector is empty or has over 255 columns.
    /// * Number of elements in any of the columns is smaller than 8, greater than the biggest
    ///   multiplicative subgroup in the field `B`, or is not a power of two.
    /// * Number of elements is not identical for all columns.
    pub fn init(columns: Vec<Vec<B>>) -> Self {
        assert!(
            !columns.is_empty(),
            "execution trace must consist of at least one column"
        );
        assert!(
            columns.len() <= TraceInfo::MAX_TRACE_WIDTH,
            "execution trace width cannot be greater than {}, but was {}",
            TraceInfo::MAX_TRACE_WIDTH,
            columns.len()
        );
        let trace_length = columns[0].len();
        assert!(
            trace_length >= TraceInfo::MIN_TRACE_LENGTH,
            "execution trace must be at least {} steps long, but was {}",
            TraceInfo::MIN_TRACE_LENGTH,
            trace_length
        );
        assert!(
            trace_length.is_power_of_two(),
            "execution trace length must be a power of 2"
        );
        assert!(
            trace_length.ilog2() <= B::TWO_ADICITY,
            "execution trace length cannot exceed 2^{} steps, but was 2^{}",
            B::TWO_ADICITY,
            trace_length.ilog2()
        );
        for column in columns.iter().skip(1) {
            assert_eq!(
                column.len(),
                trace_length,
                "all columns traces must have the same length"
            );
        }

        Self {
            layout: TraceLayout::new(columns.len(), [0], [0]),
            trace: ColMatrix::new(columns),
            meta: vec![],
        }
    }

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

    /// Updates a value in a single cell of the execution trace.
    ///
    /// Specifically, the value in the specified `column` and the specified `step` is set to the
    /// provide `value`.
    ///
    /// # Panics
    /// Panics if either `column` or `step` are out of bounds for this execution trace.
    pub fn set(&mut self, column: usize, step: usize, value: B) {
        self.trace.set(column, step, value)
    }

    /// Updates metadata for this execution trace to the specified vector of bytes.
    ///
    /// # Panics
    /// Panics if the length of `meta` is greater than 65535;
    pub fn set_meta(&mut self, meta: Vec<u8>) {
        assert!(
            meta.len() <= TraceInfo::MAX_META_LENGTH,
            "number of metadata bytes cannot be greater than {}, but was {}",
            TraceInfo::MAX_META_LENGTH,
            meta.len()
        );
        self.meta = meta
    }

    /// Fill all rows in the execution trace.
    ///
    /// The rows are filled by executing the provided closures as follows:
    /// - `init` closure is used to initialize the first row of the trace; it receives a mutable
    ///   reference to the first state initialized to all zeros. The contents of the state are
    ///   copied into the first row of the trace after the closure returns.
    /// - `update` closure is used to populate all subsequent rows of the trace; it receives two
    ///   parameters:
    ///   - index of the last updated row (starting with 0).
    ///   - a mutable reference to the last updated state; the contents of the state are copied
    ///     into the next row of the trace after the closure returns.
    pub fn fill<I, U>(&mut self, init: I, mut update: U)
    where
        I: FnOnce(&mut [B]),
        U: FnMut(usize, &mut [B]),
    {
        let mut state = vec![B::ZERO; self.main_trace_width()];
        init(&mut state);
        self.update_row(0, &state);

        for i in 0..self.length() - 1 {
            update(i, &mut state);
            self.update_row(i + 1, &state);
        }
    }

    /// Updates a single row in the execution trace with provided data.
    pub fn update_row(&mut self, step: usize, state: &[B]) {
        self.trace.update_row(step, state);
    }

    // FRAGMENTS
    // --------------------------------------------------------------------------------------------

    /// Breaks the execution trace into mutable fragments.
    ///
    /// The number of rows in each fragment will be equal to `fragment_length` parameter. The
    /// returned fragments can be used to update data in the trace from multiple threads.
    ///
    /// # Panics
    /// Panics if `fragment_length` is smaller than 2, greater than the length of the trace,
    /// or is not a power of two.
    #[cfg(not(feature = "concurrent"))]
    pub fn fragments(&mut self, fragment_length: usize) -> vec::IntoIter<TraceTableFragment<B>> {
        self.build_fragments(fragment_length).into_iter()
    }

    /// Breaks the execution trace into mutable fragments.
    ///
    /// The number of rows in each fragment will be equal to `fragment_length` parameter. The
    /// returned fragments can be used to update data in the trace from multiple threads.
    ///
    /// # Panics
    /// Panics if `fragment_length` is smaller than 2, greater than the length of the trace,
    /// or is not a power of two.
    #[cfg(feature = "concurrent")]
    pub fn fragments(
        &mut self,
        fragment_length: usize,
    ) -> rayon::vec::IntoIter<TraceTableFragment<B>> {
        self.build_fragments(fragment_length).into_par_iter()
    }

    /// Returns a vector of trace fragments each covering the number of steps specified by the
    /// `fragment_length` parameter.
    fn build_fragments(&mut self, fragment_length: usize) -> Vec<TraceTableFragment<B>> {
        assert!(
            fragment_length >= MIN_FRAGMENT_LENGTH,
            "fragment length must be at least {MIN_FRAGMENT_LENGTH}, but was {fragment_length}"
        );
        assert!(
            fragment_length <= self.length(),
            "length of a fragment cannot exceed {}, but was {}",
            self.length(),
            fragment_length
        );
        assert!(
            fragment_length.is_power_of_two(),
            "fragment length must be a power of 2"
        );
        let num_fragments = self.length() / fragment_length;

        let mut fragment_data = (0..num_fragments).map(|_| Vec::new()).collect::<Vec<_>>();
        self.trace.columns_mut().for_each(|column| {
            for (i, fragment) in column.chunks_mut(fragment_length).enumerate() {
                fragment_data[i].push(fragment);
            }
        });

        fragment_data
            .into_iter()
            .enumerate()
            .map(|(i, data)| TraceTableFragment {
                index: i,
                offset: i * fragment_length,
                data,
            })
            .collect()
    }

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

    /// Returns the number of columns in this execution trace.
    pub fn width(&self) -> usize {
        self.main_trace_width()
    }

    /// Returns the entire trace column at the specified index.
    pub fn get_column(&self, col_idx: usize) -> &[B] {
        self.trace.get_column(col_idx)
    }

    /// Returns value of the cell in the specified column at the specified row of this trace.
    pub fn get(&self, column: usize, step: usize) -> B {
        self.trace.get(column, step)
    }

    /// Reads a single row from this execution trace into the provided target.
    pub fn read_row_into(&self, step: usize, target: &mut [B]) {
        self.trace.read_row_into(step, target);
    }
}

// TRACE TRAIT IMPLEMENTATION
// ================================================================================================

impl<B: StarkField> Trace for TraceTable<B> {
    type BaseField = B;

    fn layout(&self) -> &TraceLayout {
        &self.layout
    }

    fn length(&self) -> usize {
        self.trace.num_rows()
    }

    fn meta(&self) -> &[u8] {
        &self.meta
    }

    fn read_main_frame(&self, row_idx: usize, frame: &mut EvaluationFrame<Self::BaseField>) {
        let next_row_idx = (row_idx + 1) % self.length();
        self.trace.read_row_into(row_idx, frame.current_mut());
        self.trace.read_row_into(next_row_idx, frame.next_mut());
    }

    fn main_segment(&self) -> &ColMatrix<B> {
        &self.trace
    }

    fn build_aux_segment<E>(
        &mut self,
        _aux_segments: &[ColMatrix<E>],
        _rand_elements: &[E],
    ) -> Option<ColMatrix<E>>
    where
        E: FieldElement<BaseField = Self::BaseField>,
    {
        None
    }
}

// TRACE FRAGMENTS
// ================================================================================================
/// A set of consecutive rows of an execution trace.
///
/// An execution trace fragment is a "view" into the specific execution trace. Updating data in
/// the fragment, directly updates the data in the underlying execution trace.
///
/// A fragment cannot be instantiated directly but is created by executing
/// [TraceTable::fragments()] method.
///
/// A fragment always contains contiguous rows, and the number of rows is guaranteed to be a power
/// of two.
pub struct TraceTableFragment<'a, B: StarkField> {
    index: usize,
    offset: usize,
    data: Vec<&'a mut [B]>,
}

impl<'a, B: StarkField> TraceTableFragment<'a, B> {
    // PUBLIC ACCESSORS
    // --------------------------------------------------------------------------------------------

    /// Returns the index of this fragment.
    pub fn index(&self) -> usize {
        self.index
    }

    /// Returns the step at which the fragment starts in the context of the original execution
    /// trace.
    pub fn offset(&self) -> usize {
        self.offset
    }

    /// Returns the number of rows in this execution trace fragment.
    pub fn length(&self) -> usize {
        self.data[0].len()
    }

    /// Returns the width of the fragment (same as the width of the underlying execution trace).
    pub fn width(&self) -> usize {
        self.data.len()
    }

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

    /// Fills all rows in the fragment.
    ///
    /// The rows are filled by executing the provided closures as follows:
    /// - `init` closure is used to initialize the first row of the fragment; it receives a
    ///   mutable reference to the first state initialized to all zeros. Contents of the state are
    ///   copied into the first row of the fragment after the closure returns.
    /// - `update` closure is used to populate all subsequent rows of the fragment; it receives two
    ///   parameters:
    ///   - index of the last updated row (starting with 0).
    ///   - a mutable reference to the last updated state; the contents of the state are copied
    ///     into the next row of the fragment after the closure returns.
    pub fn fill<I, T>(&mut self, init_state: I, mut update_state: T)
    where
        I: FnOnce(&mut [B]),
        T: FnMut(usize, &mut [B]),
    {
        let mut state = vec![B::ZERO; self.width()];
        init_state(&mut state);
        self.update_row(0, &state);

        for i in 0..self.length() - 1 {
            update_state(i, &mut state);
            self.update_row(i + 1, &state);
        }
    }

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


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

//! Feature-based re-export of common collection components.
//!
//! When `std` feature is enabled, this module exports collections from the Rust standard library.
//! When `alloc` feature is enabled, same collected are provided without relying on the Rust
//! standard library.

#[cfg(not(feature = "std"))]
pub use alloc::collections::{btree_map, btree_set, BTreeMap, BTreeSet};

#[cfg(not(feature = "std"))]
pub use alloc::vec::{self as vec, Vec};

#[cfg(feature = "std")]
pub use std::collections::{btree_map, btree_set, BTreeMap, BTreeSet};

#[cfg(feature = "std")]
pub use std::vec::{self as vec, Vec};


// 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::string::String;
use core::fmt;

// DESERIALIZATION ERROR
// ================================================================================================

/// Defines errors which can occur during deserialization.
#[derive(Debug, PartialEq, Eq)]
pub enum DeserializationError {
    /// Bytes in the input do not represent a valid value.
    InvalidValue(String),
    /// An end of input was reached before a valid value could be deserialized.
    UnexpectedEOF,
    /// Deserialization has finished but not all bytes have been consumed.
    UnconsumedBytes,
    /// An unknown error has occurred.
    UnknownError(String),
}

impl fmt::Display for DeserializationError {
    #[rustfmt::skip]
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidValue(err_msg) => {
                write!(f, "{err_msg}")
            }
            Self::UnexpectedEOF => {
                write!(f, "unexpected EOF")
            }
            Self::UnconsumedBytes => {
                write!(f, "not all bytes were consumed")
            }
            Self::UnknownError(err_msg) => {
                write!(f, "unknown error: {err_msg}")
            }
        }
    }
}


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

//! Components needed for parallel iterators.
//!
//! When `concurrent` feature is enabled, this module re-exports `rayon::prelude`. Otherwise,
//! this is an empty module.

#[cfg(feature = "concurrent")]
pub use rayon::{current_num_threads as rayon_num_threads, prelude::*};

/// Returns either a regular or a parallel iterator depending on whether `concurrent` feature
/// is enabled.
///
/// When `concurrent` feature is enabled, creates a parallel iterator; otherwise, creates a
/// regular iterator. Optionally, `min_length` can be used to specify the minimum length of
/// iterator to be processed in each thread.
///
/// Adapted from: <https://github.com/arkworks-rs/utils/blob/master/src/lib.rs>
#[macro_export]
macro_rules! iter {
    ($e: expr) => {{
        #[cfg(feature = "concurrent")]
        let result = $e.par_iter();

        #[cfg(not(feature = "concurrent"))]
        let result = $e.iter();

        result
    }};
    ($e: expr, $min_len: expr) => {{
        #[cfg(feature = "concurrent")]
        let result = $e.par_iter().with_min_len($min_len);

        #[cfg(not(feature = "concurrent"))]
        let result = $e.iter();

        result
    }};
}

/// Returns either a regular or a parallel mutable iterator depending on whether `concurrent`
/// feature is enabled.
///
/// When `concurrent` feature is enabled, creates a mutable parallel iterator; otherwise,
/// creates a regular mutable iterator. Optionally, `min_length` can be used to specify the
/// minimum length of iterator to be processed in each thread.
///
/// Adapted from: <https://github.com/arkworks-rs/utils/blob/master/src/lib.rs>
#[macro_export]
macro_rules! iter_mut {
    ($e: expr) => {{
        #[cfg(feature = "concurrent")]
        let result = $e.par_iter_mut();

        #[cfg(not(feature = "concurrent"))]
        let result = $e.iter_mut();

        result
    }};
    ($e: expr, $min_len: expr) => {{
        #[cfg(feature = "concurrent")]
        let result = $e.par_iter_mut().with_min_len($min_len);

        #[cfg(not(feature = "concurrent"))]
        let result = $e.iter_mut();

        result
    }};
}

/// Applies a procedure to the provided slice either in a single thread or multiple threads
/// based on whether `concurrent` feature is enabled.
///
/// When `concurrent` feature is enabled, breaks the slice into batches and processes each
/// batch in a separate thread; otherwise, the entire slice is processed as a single batch
/// in one thread. Optionally, `min_batch_size` can be used to specify the minimum size of
/// the resulting batches.
#[macro_export]
macro_rules! batch_iter_mut {
    ($e: expr, $c: expr) => {
        #[cfg(feature = "concurrent")]
        {
            let batch_size = $e.len() / rayon_num_threads().next_power_of_two();
            if batch_size < 1 {
                $c($e, 0);
            }
            else {
                $e.par_chunks_mut(batch_size).enumerate().for_each(|(i, batch)| {
                    $c(batch, i * batch_size);
                });
            }
        }

        #[cfg(not(feature = "concurrent"))]
        $c($e, 0);
    };
    ($e: expr, $min_batch_size: expr, $c: expr) => {
        #[cfg(feature = "concurrent")]
        {
            let batch_size = $e.len() / rayon_num_threads().next_power_of_two();
            if batch_size < $min_batch_size {
                $c($e, 0);
            }
            else {
                $e.par_chunks_mut(batch_size).enumerate().for_each(|(i, batch)| {
                    $c(batch, i * batch_size);
                });
            }
        }

        #[cfg(not(feature = "concurrent"))]
        $c($e, 0);
    };
}


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

//! This crate contains utility traits, functions, and macros used by other crates of Winterfell
//! STARK prover and verifier.

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

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

pub mod collections;
pub mod iterators;
pub mod string;

use collections::Vec;
use core::{convert::TryInto, mem, slice};

mod serde;
pub use serde::{ByteReader, ByteWriter, Deserializable, Serializable, SliceReader};

mod errors;
pub use errors::DeserializationError;

#[cfg(test)]
mod tests;

// FEATURE-BASED RE-EXPORTS
// ================================================================================================

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

#[cfg(feature = "concurrent")]
pub use rayon;

#[cfg(not(feature = "std"))]
pub use alloc::boxed::Box;

#[cfg(feature = "std")]
pub use std::boxed::Box;

// AS BYTES
// ================================================================================================

/// Defines a zero-copy representation of `Self` as a sequence of bytes.
pub trait AsBytes {
    /// Returns a byte representation of `self`.
    ///
    /// This method is intended to re-interpret the underlying memory as a sequence of bytes, and
    /// thus, should be zero-copy.
    fn as_bytes(&self) -> &[u8];
}

impl<const N: usize, const M: usize> AsBytes for [[u8; N]; M] {
    /// Flattens a two-dimensional array of bytes into a slice of bytes.
    fn as_bytes(&self) -> &[u8] {
        let p = self.as_ptr();
        let len = N * M;
        unsafe { slice::from_raw_parts(p as *const u8, len) }
    }
}

impl<const N: usize> AsBytes for [[u8; N]] {
    /// Flattens a slice of byte arrays into a slice of bytes.
    fn as_bytes(&self) -> &[u8] {
        let p = self.as_ptr();
        let len = self.len() * N;
        unsafe { slice::from_raw_parts(p as *const u8, len) }
    }
}

// VECTOR FUNCTIONS
// ================================================================================================

/// Returns a vector of the specified length with un-initialized memory.
///
/// This is usually faster than requesting a vector with initialized memory and is useful when we
/// overwrite all contents of the vector immediately after memory allocation.
///
/// # Safety
/// Using values from the returned vector before initializing them will lead to undefined behavior.
#[allow(clippy::uninit_vec)]
pub unsafe fn uninit_vector<T>(length: usize) -> Vec<T> {
    let mut vector = Vec::with_capacity(length);
    vector.set_len(length);
    vector
}

// GROUPING / UN-GROUPING FUNCTIONS
// ================================================================================================

/// Transmutes a vector of `n` elements into a vector of `n` / `N` elements, each of which is
/// an array of `N` elements.
///
/// This function just re-interprets the underlying memory and is thus zero-copy.
/// # Panics
/// Panics if `n` is not divisible by `N`.
///
/// # Example
/// ```
/// # use winter_utils::group_vector_elements;
/// let a = vec![0_u32, 1, 2, 3, 4, 5, 6, 7];
/// let b: Vec<[u32; 2]> = group_vector_elements(a);
///
/// assert_eq!(vec![[0, 1], [2, 3], [4, 5], [6, 7]], b);
/// ```
pub fn group_vector_elements<T, const N: usize>(source: Vec<T>) -> Vec<[T; N]> {
    assert_eq!(
        source.len() % N,
        0,
        "source length must be divisible by {}, but was {}",
        N,
        source.len()
    );
    let mut v = mem::ManuallyDrop::new(source);
    let p = v.as_mut_ptr();
    let len = v.len() / N;
    let cap = v.capacity() / N;
    unsafe { Vec::from_raw_parts(p as *mut [T; N], len, cap) }
}

/// Transmutes a slice of `n` elements into a slice of `n` / `N` elements, each of which is
/// an array of `N` elements.
///
/// This function just re-interprets the underlying memory and is thus zero-copy.
/// # Panics
/// Panics if `n` is not divisible by `N`.
///
/// # Example
/// ```
/// # use winter_utils::group_slice_elements;
/// let a = [0_u32, 1, 2, 3, 4, 5, 6, 7];
/// let b: &[[u32; 2]] = group_slice_elements(&a);
///
/// assert_eq!(&[[0, 1], [2, 3], [4, 5], [6, 7]], b);
/// ```
pub fn group_slice_elements<T, const N: usize>(source: &[T]) -> &[[T; N]] {
    assert_eq!(
        source.len() % N,
        0,
        "source length must be divisible by {N}"
    );
    let p = source.as_ptr();
    let len = source.len() / N;
    unsafe { slice::from_raw_parts(p as *const [T; N], len) }
}

/// Transmutes a slice of `n` arrays each of length `N`, into a slice of `N` * `n` elements.
///
/// This function just re-interprets the underlying memory and is thus zero-copy.
/// # Example
/// ```
/// # use winter_utils::flatten_slice_elements;
/// let a = vec![[1, 2, 3, 4], [5, 6, 7, 8]];
///
/// let b = flatten_slice_elements(&a);
/// assert_eq!(&[1, 2, 3, 4, 5, 6, 7, 8], b);
/// ```
pub fn flatten_slice_elements<T, const N: usize>(source: &[[T; N]]) -> &[T] {
    let p = source.as_ptr();
    let len = source.len() * N;
    unsafe { slice::from_raw_parts(p as *const T, len) }
}

/// Transmutes a vector of `n` arrays each of length `N`, into a vector of `N` * `n` elements.
///
/// This function just re-interprets the underlying memory and is thus zero-copy.
/// # Example
/// ```
/// # use winter_utils::flatten_vector_elements;
/// let a = vec![[1, 2, 3, 4], [5, 6, 7, 8]];
///
/// let b = flatten_vector_elements(a);
/// assert_eq!(vec![1, 2, 3, 4, 5, 6, 7, 8], b);
/// ```
pub fn flatten_vector_elements<T, const N: usize>(source: Vec<[T; N]>) -> Vec<T> {
    let v = mem::ManuallyDrop::new(source);
    let p = v.as_ptr();
    let len = v.len() * N;
    let cap = v.capacity() * N;
    unsafe { Vec::from_raw_parts(p as *mut T, len, cap) }
}

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

/// Transposes a slice of `n` elements into a matrix with `N` columns and `n`/`N` rows.
///
/// When `concurrent` feature is enabled, the slice will be transposed using multiple threads.
///
/// # Panics
/// Panics if `n` is not divisible by `N`.
///
/// # Example
/// ```
/// # use winter_utils::transpose_slice;
/// let a = [0_u32, 1, 2, 3, 4, 5, 6, 7];
/// let b: Vec<[u32; 2]> = transpose_slice(&a);
///
/// assert_eq!(vec![[0, 4], [1, 5], [2, 6], [3, 7]], b);
/// ```
pub fn transpose_slice<T: Copy + Send + Sync, const N: usize>(source: &[T]) -> Vec<[T; N]> {
    let row_count = source.len() / N;
    assert_eq!(
        row_count * N,
        source.len(),
        "source length must be divisible by {}, but was {}",
        N,
        source.len()
    );

    let mut result = unsafe { group_vector_elements(uninit_vector(row_count * N)) };
    iter_mut!(result, 1024)
        .enumerate()
        .for_each(|(i, element)| {
            for j in 0..N {
                element[j] = source[i + j * row_count]
            }
        });
    result
}

// RANDOMNESS
// ================================================================================================

/// Defines how `Self` can be read from a sequence of random bytes.
pub trait Randomizable: Sized {
    /// Size of `Self` in bytes.
    ///
    /// This is used to determine how many bytes should be passed to the
    /// [from_random_bytes()](Self::from_random_bytes) function.
    const VALUE_SIZE: usize;

    /// Returns `Self` if the set of bytes forms a valid value, otherwise returns None.
    fn from_random_bytes(source: &[u8]) -> Option<Self>;
}

impl Randomizable for u128 {
    const VALUE_SIZE: usize = 16;

    fn from_random_bytes(source: &[u8]) -> Option<Self> {
        if let Ok(bytes) = source[..Self::VALUE_SIZE].try_into() {
            Some(u128::from_le_bytes(bytes))
        } else {
            None
        }
    }
}

impl Randomizable for u64 {
    const VALUE_SIZE: usize = 8;

    fn from_random_bytes(source: &[u8]) -> Option<Self> {
        if let Ok(bytes) = source[..Self::VALUE_SIZE].try_into() {
            Some(u64::from_le_bytes(bytes))
        } else {
            None
        }
    }
}

impl Randomizable for u32 {
    const VALUE_SIZE: usize = 4;

    fn from_random_bytes(source: &[u8]) -> Option<Self> {
        if let Ok(bytes) = source[..Self::VALUE_SIZE].try_into() {
            Some(u32::from_le_bytes(bytes))
        } else {
            None
        }
    }
}

impl Randomizable for u16 {
    const VALUE_SIZE: usize = 2;

    fn from_random_bytes(source: &[u8]) -> Option<Self> {
        if let Ok(bytes) = source[..Self::VALUE_SIZE].try_into() {
            Some(u16::from_le_bytes(bytes))
        } else {
            None
        }
    }
}

impl Randomizable for u8 {
    const VALUE_SIZE: usize = 1;

    fn from_random_bytes(source: &[u8]) -> Option<Self> {
        Some(source[0])
    }
}


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

//! Feature-based re-export of common string components.
//!
//! When `std` feature is enabled, this module exports string components from the Rust standard
//! library. When `alloc` feature is enabled, same components are provided without relying on the
//! Rust standard library.

#[cfg(not(feature = "std"))]
pub use alloc::string::{String, ToString};

#[cfg(feature = "std")]
pub use std::string::{String, ToString};


// 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::{collections::Vec, ByteReader, ByteWriter, Serializable, SliceReader};

// VECTOR UTILS TESTS
// ================================================================================================

#[test]
fn group_vector_elements() {
    let n = 16;
    let a = (0..n).map(|v| v as u64).collect::<Vec<_>>();

    let b = super::group_vector_elements::<u64, 4>(a.clone());
    for i in 0..b.len() {
        for j in 0..4 {
            assert_eq!(a[i * 4 + j], b[i][j]);
        }
    }

    let b = super::group_vector_elements::<u64, 2>(a.clone());
    for i in 0..b.len() {
        for j in 0..2 {
            assert_eq!(a[i * 2 + j], b[i][j]);
        }
    }
}

// SLICE READER TESTS
// ================================================================================================

#[test]
fn read_u8() {
    let source = [1u8, 3, 5, 7];
    let mut a = SliceReader::new(&source);

    assert_eq!(1, a.read_u8().unwrap());
    assert_eq!(3, a.read_u8().unwrap());
    assert_eq!(5, a.read_u8().unwrap());
    assert_eq!(7, a.read_u8().unwrap());
    assert!(a.read_u8().is_err());
}

#[test]
fn read_u16() {
    let mut source = 12345u16.to_le_bytes().to_vec();
    source.append(&mut 23456u16.to_le_bytes().to_vec());
    let mut a = SliceReader::new(&source);

    assert_eq!(12345, a.read_u16().unwrap());
    assert_eq!(23456, a.read_u16().unwrap());
    assert!(a.read_u16().is_err());
}

#[test]
fn read_u32() {
    let mut source = 123456789u32.to_le_bytes().to_vec();
    source.append(&mut 2345678910u32.to_le_bytes().to_vec());
    let mut a = SliceReader::new(&source);

    assert_eq!(123456789, a.read_u32().unwrap());
    assert_eq!(2345678910, a.read_u32().unwrap());
    assert!(a.read_u32().is_err());
}

#[test]
fn read_u64() {
    let mut source = 12345678910u64.to_le_bytes().to_vec();
    source.append(&mut 234567891011u64.to_le_bytes().to_vec());
    let mut a = SliceReader::new(&source);

    assert_eq!(12345678910, a.read_u64().unwrap());
    assert_eq!(234567891011, a.read_u64().unwrap());
    assert!(a.read_u64().is_err());
}

#[test]
fn read_u8_vec() {
    let source = [1u8, 2, 3, 4, 5, 6, 7, 8];
    let mut a = SliceReader::new(&source);

    assert_eq!(vec![1, 2], a.read_vec(2).unwrap());
    assert_eq!(vec![3, 4, 5], a.read_vec(3).unwrap());
    assert_eq!(vec![6, 7], a.read_vec(2).unwrap());
    assert_eq!(vec![8], a.read_vec(1).unwrap());
    assert!(a.read_vec(2).is_err());
}

// SERIALIZATION TESTS
// ================================================================================================

impl Serializable for u128 {
    fn write_into<W: crate::ByteWriter>(&self, target: &mut W) {
        target.write_bytes(&self.to_le_bytes());
    }
}

#[test]
fn write_serializable() {
    let mut target: Vec<u8> = Vec::new();

    123456u128.write_into(&mut target);
    assert_eq!(16, target.len());

    target.write(234567u128);
    assert_eq!(32, target.len());

    let mut reader = SliceReader::new(&target);
    assert_eq!(123456u128, reader.read_u128().unwrap());
    assert_eq!(234567u128, reader.read_u128().unwrap());
}

#[test]
fn write_serializable_batch() {
    let mut target: Vec<u8> = Vec::new();

    let batch1 = vec![1u128, 2, 3, 4];
    batch1.write_into(&mut target);
    assert_eq!(64, target.len());

    let batch2 = [5u128, 6, 7, 8];
    target.write(&batch2[..]);
    assert_eq!(128, target.len());

    let mut reader = SliceReader::new(&target);
    for i in 1u128..9 {
        assert_eq!(i, reader.read_u128().unwrap());
    }
}

#[test]
fn write_serializable_array_batch() {
    let mut target: Vec<u8> = Vec::new();

    let batch1 = vec![[1u128, 2], [3, 4]];
    batch1.write_into(&mut target);
    assert_eq!(64, target.len());

    let batch2 = [[5u128, 6], [7, 8]];
    target.write(&batch2[..]);
    assert_eq!(128, target.len());

    let mut reader = SliceReader::new(&target);
    for i in 1u128..9 {
        assert_eq!(i, reader.read_u128().unwrap());
    }
}


// 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::{Deserializable, DeserializationError, Vec};

// BYTE READER TRAIT
// ================================================================================================

/// Defines how primitive values are to be read from `Self`.
pub trait ByteReader {
    // REQUIRED METHODS
    // --------------------------------------------------------------------------------------------

    /// Returns a single byte read from `self`.
    ///
    /// # Errors
    /// Returns a [DeserializationError] error the reader is at EOF.
    fn read_u8(&mut self) -> Result<u8, DeserializationError>;

    /// Returns the next byte to be read from `self` without advancing the reader to the next byte.
    ///
    /// # Errors
    /// Returns a [DeserializationError] error the reader is at EOF.
    fn peek_u8(&self) -> Result<u8, DeserializationError>;

    /// Returns a byte vector of the specified length read from `self`.
    ///
    /// # Errors
    /// Returns a [DeserializationError] if a vector of the specified length could not be read
    /// from `self`.
    fn read_vec(&mut self, len: usize) -> Result<Vec<u8>, DeserializationError>;

    /// Returns a byte array of length `N` read from `self`.
    ///
    /// # Errors
    /// Returns a [DeserializationError] if an array of the specified length could not be read
    /// from `self`.
    fn read_array<const N: usize>(&mut self) -> Result<[u8; N], DeserializationError>;

    /// Checks if it is possible to read at least `num_bytes` bytes from this ByteReader
    ///
    /// # Errors
    /// Returns an error if, when reading the requested number of bytes, we go beyond the
    /// boundaries of the array
    fn check_eor(&self, num_bytes: usize) -> Result<(), DeserializationError>;

    /// Returns true if there are more bytes left to be read from `self`.
    fn has_more_bytes(&self) -> bool;

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

    /// Returns a boolean value read from `self` consuming 1 byte from the reader.
    ///
    /// # Errors
    /// Returns a [DeserializationError] if a u16 value could not be read from `self`.
    fn read_bool(&mut self) -> Result<bool, DeserializationError> {
        let byte = self.read_u8()?;
        match byte {
            0 => Ok(false),
            1 => Ok(true),
            _ => Err(DeserializationError::InvalidValue(format!(
                "{byte} is not a boolean value"
            ))),
        }
    }

    /// Returns a u16 value read from `self` in little-endian byte order.
    ///
    /// # Errors
    /// Returns a [DeserializationError] if a u16 value could not be read from `self`.
    fn read_u16(&mut self) -> Result<u16, DeserializationError> {
        let bytes = self.read_array::<2>()?;
        Ok(u16::from_le_bytes(bytes))
    }

    /// Returns a u32 value read from `self` in little-endian byte order.
    ///
    /// # Errors
    /// Returns a [DeserializationError] if a u32 value could not be read from `self`.
    fn read_u32(&mut self) -> Result<u32, DeserializationError> {
        let bytes = self.read_array::<4>()?;
        Ok(u32::from_le_bytes(bytes))
    }

    /// Returns a u64 value read from `self` in little-endian byte order.
    ///
    /// # Errors
    /// Returns a [DeserializationError] if a u64 value could not be read from `self`.
    fn read_u64(&mut self) -> Result<u64, DeserializationError> {
        let bytes = self.read_array::<8>()?;
        Ok(u64::from_le_bytes(bytes))
    }

    /// Returns a u128 value read from `self` in little-endian byte order.
    ///
    /// # Errors
    /// Returns a [DeserializationError] if a u128 value could not be read from `self`.
    fn read_u128(&mut self) -> Result<u128, DeserializationError> {
        let bytes = self.read_array::<16>()?;
        Ok(u128::from_le_bytes(bytes))
    }

    /// Reads a deserializable value from `self`.
    ///
    /// # Panics
    /// Panics if the value could not be read from `self`.
    fn read<D>(&mut self) -> Result<D, DeserializationError>
    where
        Self: Sized,
        D: Deserializable,
    {
        D::read_from(self)
    }
}

// SLICE READER
// ================================================================================================

/// Implements [ByteReader] trait for a slice of bytes.
pub struct SliceReader<'a> {
    source: &'a [u8],
    pos: usize,
}

impl<'a> SliceReader<'a> {
    /// Creates a new slice reader from the specified slice.
    pub fn new(source: &'a [u8]) -> Self {
        SliceReader { source, pos: 0 }
    }
}

impl<'a> ByteReader for SliceReader<'a> {
    fn read_u8(&mut self) -> Result<u8, DeserializationError> {
        self.check_eor(1)?;
        let result = self.source[self.pos];
        self.pos += 1;
        Ok(result)
    }

    fn peek_u8(&self) -> Result<u8, DeserializationError> {
        self.check_eor(1)?;
        Ok(self.source[self.pos])
    }

    fn read_vec(&mut self, len: usize) -> Result<Vec<u8>, DeserializationError> {
        self.check_eor(len)?;
        let result = self.source[self.pos..self.pos + len].to_vec();
        self.pos += len;
        Ok(result)
    }

    fn read_array<const N: usize>(&mut self) -> Result<[u8; N], DeserializationError> {
        self.check_eor(N)?;
        let mut result = [0_u8; N];
        result.copy_from_slice(&self.source[self.pos..self.pos + N]);
        self.pos += N;
        Ok(result)
    }

    fn check_eor(&self, num_bytes: usize) -> Result<(), DeserializationError> {
        if self.pos + num_bytes > self.source.len() {
            return Err(DeserializationError::UnexpectedEOF);
        }
        Ok(())
    }

    fn has_more_bytes(&self) -> bool {
        self.pos < self.source.len()
    }
}


// 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::{Serializable, Vec};

// BYTE WRITER TRAIT
// ================================================================================================

/// Defines how primitive values are to be written into `Self`.
pub trait ByteWriter: Sized {
    // REQUIRED METHODS
    // --------------------------------------------------------------------------------------------

    /// Writes a single byte into `self`.
    ///
    /// # Panics
    /// Panics if the byte could not be written into `self`.
    fn write_u8(&mut self, value: u8);

    /// Writes a sequence of bytes into `self`.
    ///
    /// # Panics
    /// Panics if the sequence of bytes could not be written into `self`.
    fn write_bytes(&mut self, values: &[u8]);

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

    /// Writes a boolean value into `self`.
    ///
    /// A boolean value is written as a single byte.
    ///
    /// # Panics
    /// Panics if the value could not be written into `self`.
    fn write_bool(&mut self, val: bool) {
        self.write_u8(val as u8);
    }

    /// Writes a u16 value in little-endian byte order into `self`.
    ///
    /// # Panics
    /// Panics if the value could not be written into `self`.
    fn write_u16(&mut self, value: u16) {
        self.write_bytes(&value.to_le_bytes());
    }

    /// Writes a u32 value in little-endian byte order into `self`.
    ///
    /// # Panics
    /// Panics if the value could not be written into `self`.
    fn write_u32(&mut self, value: u32) {
        self.write_bytes(&value.to_le_bytes());
    }

    /// Writes a u64 value in little-endian byte order into `self`.
    ///
    /// # Panics
    /// Panics if the value could not be written into `self`.
    fn write_u64(&mut self, value: u64) {
        self.write_bytes(&value.to_le_bytes());
    }

    /// Writes a serializable value into `self`.
    ///
    /// # Panics
    /// Panics if the value could not be written into `self`.
    fn write<S: Serializable>(&mut self, value: S) {
        value.write_into(self)
    }
}

// BYTE WRITER IMPLEMENTATIONS
// ================================================================================================

impl ByteWriter for Vec<u8> {
    fn write_u8(&mut self, value: u8) {
        self.push(value);
    }

    fn write_bytes(&mut self, values: &[u8]) {
        self.extend_from_slice(values);
    }
}