// 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::{
    rescue::{self, STATE_WIDTH},
    BaseElement, ExtensionOf, FieldElement, ProofOptions, CYCLE_LENGTH, TRACE_WIDTH,
};
use crate::utils::{are_equal, not, EvaluationResult};
use core_utils::flatten_slice_elements;
use winterfell::{
    math::ToElements, Air, AirContext, Assertion, AuxTraceRandElements, EvaluationFrame, TraceInfo,
    TransitionConstraintDegree,
};

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

/// Specifies steps on which Rescue transition function is applied.
const CYCLE_MASK: [BaseElement; CYCLE_LENGTH] = [
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ZERO,
    BaseElement::ZERO,
];

// RESCUE AIR
// ================================================================================================

pub struct PublicInputs {
    pub result: [[BaseElement; 2]; 2],
}

impl ToElements<BaseElement> for PublicInputs {
    fn to_elements(&self) -> Vec<BaseElement> {
        flatten_slice_elements(&self.result).to_vec()
    }
}

pub struct RescueRapsAir {
    context: AirContext<BaseElement>,
    result: [[BaseElement; 2]; 2],
}

impl Air for RescueRapsAir {
    type BaseField = BaseElement;
    type PublicInputs = PublicInputs;

    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    fn new(trace_info: TraceInfo, pub_inputs: PublicInputs, options: ProofOptions) -> Self {
        let main_degrees =
            vec![TransitionConstraintDegree::with_cycles(3, vec![CYCLE_LENGTH]); 2 * STATE_WIDTH];
        let aux_degrees = vec![
            TransitionConstraintDegree::with_cycles(1, vec![CYCLE_LENGTH]),
            TransitionConstraintDegree::with_cycles(1, vec![CYCLE_LENGTH]),
            TransitionConstraintDegree::new(2),
        ];
        assert_eq!(TRACE_WIDTH + 3, trace_info.width());
        RescueRapsAir {
            context: AirContext::new_multi_segment(
                trace_info,
                main_degrees,
                aux_degrees,
                8,
                2,
                options,
            ),
            result: pub_inputs.result,
        }
    }

    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],
    ) {
        let current = frame.current();
        let next = frame.next();
        // expected state width is 2*4 field elements
        debug_assert_eq!(TRACE_WIDTH, current.len());
        debug_assert_eq!(TRACE_WIDTH, next.len());

        // split periodic values into hash_flag, absorption flag and Rescue round constants
        let hash_flag = periodic_values[0];
        let absorption_flag = periodic_values[1];
        let ark = &periodic_values[2..];

        // when hash_flag = 1, constraints for Rescue round are enforced (steps 0 to 14)
        // Enforcing the round for the first hash chain
        rescue::enforce_round(
            &mut result[..STATE_WIDTH],
            &current[..STATE_WIDTH],
            &next[..STATE_WIDTH],
            ark,
            hash_flag,
        );

        // Enforcing the round for the second hash chain
        rescue::enforce_round(
            &mut result[STATE_WIDTH..],
            &current[STATE_WIDTH..],
            &next[STATE_WIDTH..],
            ark,
            hash_flag,
        );

        // When absorbing the additional seeds (step 14), we do not verify correctness of the
        // rate registers. Instead, we only verify that capacity registers have not
        // changed. When computing the permutation argument, we will recompute the permuted
        // values from the contiguous rows.
        // At step 15, we enforce that the whole hash states are copied to the next step,
        // enforcing that the values added to the capacity registers at step 14 and used in the
        // permutation argument are the ones being used in the next hashing sequence.
        result.agg_constraint(2, absorption_flag, are_equal(current[2], next[2]));
        result.agg_constraint(3, absorption_flag, are_equal(current[3], next[3]));

        result.agg_constraint(6, absorption_flag, are_equal(current[6], next[6]));
        result.agg_constraint(7, absorption_flag, are_equal(current[7], next[7]));

        // when hash_flag + absorption_flag = 0, constraints for copying hash values to the
        // next step are enforced.
        let copy_flag = not(hash_flag + absorption_flag);
        enforce_hash_copy(
            &mut result[..STATE_WIDTH],
            &current[..STATE_WIDTH],
            &next[..STATE_WIDTH],
            copy_flag,
        );
        enforce_hash_copy(
            &mut result[STATE_WIDTH..],
            &current[STATE_WIDTH..],
            &next[STATE_WIDTH..],
            copy_flag,
        );
    }

    fn evaluate_aux_transition<F, E>(
        &self,
        main_frame: &EvaluationFrame<F>,
        aux_frame: &EvaluationFrame<E>,
        periodic_values: &[F],
        aux_rand_elements: &AuxTraceRandElements<E>,
        result: &mut [E],
    ) where
        F: FieldElement<BaseField = Self::BaseField>,
        E: FieldElement<BaseField = Self::BaseField> + ExtensionOf<F>,
    {
        let main_current = main_frame.current();
        let main_next = main_frame.next();

        let aux_current = aux_frame.current();
        let aux_next = aux_frame.next();

        let random_elements = aux_rand_elements.get_segment_elements(0);

        let absorption_flag = periodic_values[1];

        // We want to enforce that the absorbed values of the first hash chain are a
        // permutation of the absorbed values of the second one. Recall that the type
        // for both seed and permuted_seed (the arrays being hashed into the chain), was
        // [[BaseElement; 2]] and we never permute any of the internal arrays, since
        // each [BaseElement; 2] represents the capacity registers for a single link in the
        // hash chain. Due to this, we want to copy two values per hash chain at iteration
        // (namely, the two capacity registers). To reduce the number of auxiliary registers needed
        // to represent each link, we group them with random elements into a single cell via
        // Î±_0 * c_0 + Î±_1 * c_1, where c_i is computed as next_i - current_i.

        // Note that the reason we use next_i - current_i is that we are
        // absorbing the new seed by adding it to the output of the previous hash.

        // Note that storing the copied values into two auxiliary columns. One could
        // instead directly compute the permutation argument, hence require a single
        // auxiliary one. For the sake of illustrating RAPs behaviour, we will store
        // the computed values in additional columns.

        let copied_value_1 = random_elements[0] * (main_next[0] - main_current[0]).into()
            + random_elements[1] * (main_next[1] - main_current[1]).into();

        result.agg_constraint(
            0,
            absorption_flag.into(),
            are_equal(aux_current[0], copied_value_1),
        );

        let copied_value_2 = random_elements[0] * (main_next[4] - main_current[4]).into()
            + random_elements[1] * (main_next[5] - main_current[5]).into();

        result.agg_constraint(
            1,
            absorption_flag.into(),
            are_equal(aux_current[1], copied_value_2),
        );

        // Enforce that the permutation argument column scales at each step by (aux[0] + Î³) / (aux[1] + Î³).
        result.agg_constraint(
            2,
            E::ONE,
            are_equal(
                aux_next[2] * (aux_current[1] + random_elements[2]),
                aux_current[2] * (aux_current[0] + random_elements[2]),
            ),
        );
    }

    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
        // Assert starting and ending values of the hash chain
        let last_step = self.trace_length() - 1;
        vec![
            // Initial capacity registers must be set to zero
            Assertion::single(2, 0, BaseElement::ZERO),
            Assertion::single(3, 0, BaseElement::ZERO),
            Assertion::single(6, 0, BaseElement::ZERO),
            Assertion::single(7, 0, BaseElement::ZERO),
            // Final rate registers (digests) should be equal to
            // the provided public input
            Assertion::single(0, last_step, self.result[0][0]),
            Assertion::single(1, last_step, self.result[0][1]),
            Assertion::single(4, last_step, self.result[1][0]),
            Assertion::single(5, last_step, self.result[1][1]),
        ]
    }

    fn get_aux_assertions<E: FieldElement + From<Self::BaseField>>(
        &self,
        _aux_rand_elements: &AuxTraceRandElements<E>,
    ) -> Vec<Assertion<E>> {
        let last_step = self.trace_length() - 1;
        vec![
            Assertion::single(2, 0, E::ONE),
            Assertion::single(2, last_step, E::ONE),
        ]
    }

    fn get_periodic_column_values(&self) -> Vec<Vec<Self::BaseField>> {
        let mut result = vec![CYCLE_MASK.to_vec()];
        let mut absorption_column = vec![BaseElement::ZERO; CYCLE_LENGTH];
        absorption_column[14] = BaseElement::ONE;
        result.push(absorption_column);

        result.append(&mut rescue::get_round_constants());

        result
    }
}

// HELPER EVALUATORS
// ------------------------------------------------------------------------------------------------

/// when flag = 1, enforces that the next state of the computation is defined like so:
/// - the first two registers are equal to the values from the previous step
/// - the other two registers are not restrained, they could be arbitrary elements,
///   until the RAP columns enforces that they are a permutation of the two final registers
///   of the other parallel chain
fn enforce_hash_copy<E: FieldElement>(result: &mut [E], current: &[E], next: &[E], flag: E) {
    result.agg_constraint(0, flag, are_equal(current[0], next[0]));
    result.agg_constraint(1, flag, are_equal(current[1], next[1]));
    result.agg_constraint(2, flag, are_equal(current[2], next[2]));
    result.agg_constraint(3, flag, are_equal(current[3], next[3]));
}


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

use core_utils::{collections::Vec, uninit_vector};
use winterfell::{
    math::{FieldElement, StarkField},
    ColMatrix, EvaluationFrame, Trace, TraceInfo, TraceLayout,
};

// RAP TRACE TABLE
// ================================================================================================
/// A concrete implementation of the [Trace] trait supporting custom RAPs.
///
/// This implementation should be sufficient for most use cases.
/// To create a trace table, you can use [RapTraceTable::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()](RapTraceTable::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 [RapTraceTable::with_meta()] function to create a blank execution trace.
/// This function work just like [RapTraceTable::new()] function, but also takes a metadata
/// parameter which can be an arbitrary sequence of bytes up to 64KB in size.
pub struct RapTraceTable<B: StarkField> {
    layout: TraceLayout,
    trace: ColMatrix<B>,
    meta: Vec<u8>,
}

impl<B: StarkField> RapTraceTable<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, [3], [3]),
            trace: ColMatrix::new(columns),
            meta,
        }
    }

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

    /// 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, update: U)
    where
        I: Fn(&mut [B]),
        U: Fn(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);
    }

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

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

    /// 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 RapTraceTable<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>,
    {
        // We only have one auxiliary segment for this example
        if !aux_segments.is_empty() {
            return None;
        }

        let mut current_row = unsafe { uninit_vector(self.width()) };
        let mut next_row = unsafe { uninit_vector(self.width()) };
        self.read_row_into(0, &mut current_row);
        let mut aux_columns = vec![vec![E::ZERO; self.length()]; self.aux_trace_width()];

        // Columns storing the copied values for the permutation argument are not necessary, but
        // help understanding the construction of RAPs and are kept for illustrative purposes.
        aux_columns[0][0] =
            rand_elements[0] * current_row[0].into() + rand_elements[1] * current_row[1].into();
        aux_columns[1][0] =
            rand_elements[0] * current_row[4].into() + rand_elements[1] * current_row[5].into();

        // Permutation argument column
        aux_columns[2][0] = E::ONE;

        for index in 1..self.length() {
            // At every last step before a new hash iteration,
            // copy the permuted values into the auxiliary columns
            if (index % super::CYCLE_LENGTH) == super::NUM_HASH_ROUNDS {
                self.read_row_into(index, &mut current_row);
                self.read_row_into(index + 1, &mut next_row);

                aux_columns[0][index] = rand_elements[0] * (next_row[0] - current_row[0]).into()
                    + rand_elements[1] * (next_row[1] - current_row[1]).into();
                aux_columns[1][index] = rand_elements[0] * (next_row[4] - current_row[4]).into()
                    + rand_elements[1] * (next_row[5] - current_row[5]).into();
            }

            let num = aux_columns[0][index - 1] + rand_elements[2];
            let denom = aux_columns[1][index - 1] + rand_elements[2];
            aux_columns[2][index] = aux_columns[2][index - 1] * num * denom.inv();
        }

        Some(ColMatrix::new(aux_columns))
    }
}


// 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::{Blake3_192, Blake3_256, Example, ExampleOptions, HashFunction, Sha3_256};
use core::marker::PhantomData;
use log::debug;
use rand_utils::rand_array;
use std::time::Instant;
use winterfell::{
    crypto::{DefaultRandomCoin, ElementHasher},
    math::{fields::f128::BaseElement, ExtensionOf, FieldElement},
    ProofOptions, Prover, StarkProof, Trace, VerifierError,
};

mod custom_trace_table;
pub use custom_trace_table::RapTraceTable;

use super::rescue::rescue::{self, STATE_WIDTH};

mod air;
use air::{PublicInputs, RescueRapsAir};

mod prover;
use prover::RescueRapsProver;

#[cfg(test)]
mod tests;

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

const CYCLE_LENGTH: usize = 16;
const NUM_HASH_ROUNDS: usize = 14;
const TRACE_WIDTH: usize = 4 * 2;

// RESCUE SPLIT HASH CHAIN EXAMPLE
// ================================================================================================

pub fn get_example(
    options: &ExampleOptions,
    chain_length: usize,
) -> Result<Box<dyn Example>, String> {
    let (options, hash_fn) = options.to_proof_options(42, 4);

    match hash_fn {
        HashFunction::Blake3_192 => Ok(Box::new(RescueRapsExample::<Blake3_192>::new(
            chain_length,
            options,
        ))),
        HashFunction::Blake3_256 => Ok(Box::new(RescueRapsExample::<Blake3_256>::new(
            chain_length,
            options,
        ))),
        HashFunction::Sha3_256 => Ok(Box::new(RescueRapsExample::<Sha3_256>::new(
            chain_length,
            options,
        ))),
        _ => Err("The specified hash function cannot be used with this example.".to_string()),
    }
}

pub struct RescueRapsExample<H: ElementHasher> {
    options: ProofOptions,
    chain_length: usize,
    seeds: Vec<[BaseElement; 2]>,
    permuted_seeds: Vec<[BaseElement; 2]>,
    result: [[BaseElement; 2]; 2],
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> RescueRapsExample<H> {
    pub fn new(chain_length: usize, options: ProofOptions) -> Self {
        assert!(
            chain_length.is_power_of_two(),
            "chain length must a power of 2"
        );
        assert!(chain_length > 2, "chain length must be at least 4");

        let mut seeds = vec![[BaseElement::ZERO; 2]; chain_length];
        for internal_seed in seeds.iter_mut() {
            *internal_seed = rand_array();
        }
        let mut permuted_seeds = seeds[2..].to_vec();
        permuted_seeds.push(seeds[0]);
        permuted_seeds.push(seeds[1]);

        // compute the sequence of hashes using external implementation of Rescue hash
        let now = Instant::now();
        let result = compute_permuted_hash_chains(&seeds, &permuted_seeds);
        debug!(
            "Computed two permuted chains of {} Rescue hashes in {} ms",
            chain_length,
            now.elapsed().as_millis(),
        );

        RescueRapsExample {
            options,
            chain_length,
            seeds,
            permuted_seeds,
            result,
            _hasher: PhantomData,
        }
    }
}

// EXAMPLE IMPLEMENTATION
// ================================================================================================

impl<H: ElementHasher> Example for RescueRapsExample<H>
where
    H: ElementHasher<BaseField = BaseElement>,
{
    fn prove(&self) -> StarkProof {
        // generate the execution trace
        debug!(
            "Generating proof for computing a chain of {} Rescue hashes\n\
            ---------------------",
            self.chain_length
        );

        // create a prover
        let prover = RescueRapsProver::<H>::new(self.options.clone());

        // generate the execution trace
        let now = Instant::now();
        let trace = prover.build_trace(&self.seeds, &self.permuted_seeds, self.result);
        let trace_length = trace.length();
        debug!(
            "Generated execution trace of {} registers and 2^{} steps in {} ms",
            trace.width(),
            trace_length.ilog2(),
            now.elapsed().as_millis()
        );

        // generate the proof
        prover.prove(trace).unwrap()
    }

    fn verify(&self, proof: StarkProof) -> Result<(), VerifierError> {
        let pub_inputs = PublicInputs {
            result: self.result,
        };
        winterfell::verify::<RescueRapsAir, H, DefaultRandomCoin<H>>(proof, pub_inputs)
    }

    fn verify_with_wrong_inputs(&self, proof: StarkProof) -> Result<(), VerifierError> {
        let pub_inputs = PublicInputs {
            result: [self.result[1], self.result[0]],
        };
        winterfell::verify::<RescueRapsAir, H, DefaultRandomCoin<H>>(proof, pub_inputs)
    }
}

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

fn absorb(state: &mut [BaseElement; STATE_WIDTH], values: &[BaseElement; 2]) {
    state[0] += values[0];
    state[1] += values[1];
    for i in 0..NUM_HASH_ROUNDS {
        rescue::apply_round(state, i);
    }
}

fn compute_permuted_hash_chains(
    seeds: &[[BaseElement; 2]],
    permuted_seeds: &[[BaseElement; 2]],
) -> [[BaseElement; 2]; 2] {
    let mut state = [BaseElement::ZERO; STATE_WIDTH];
    let mut permuted_state = [BaseElement::ZERO; STATE_WIDTH];

    // Start the hash chain
    for (seed, permuted_seed) in seeds.iter().zip(permuted_seeds) {
        absorb(&mut state, seed);
        absorb(&mut permuted_state, permuted_seed);
    }

    [[state[0], state[1]], [permuted_state[0], permuted_state[1]]]
}

fn apply_rescue_round_parallel(multi_state: &mut [BaseElement], step: usize) {
    debug_assert_eq!(multi_state.len() % STATE_WIDTH, 0);

    for state in multi_state.chunks_mut(STATE_WIDTH) {
        rescue::apply_round(state, step)
    }
}


// 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::{
    apply_rescue_round_parallel, rescue::STATE_WIDTH, BaseElement, DefaultRandomCoin,
    ElementHasher, FieldElement, PhantomData, ProofOptions, Prover, PublicInputs, RapTraceTable,
    RescueRapsAir, Trace, CYCLE_LENGTH, NUM_HASH_ROUNDS,
};

// RESCUE PROVER
// ================================================================================================
/// This example constructs a proof for correct execution of 2 hash chains simultaneously.
/// In order to demonstrate the power of RAPs, the two hash chains have seeds that are
/// permutations of each other.
pub struct RescueRapsProver<H: ElementHasher> {
    options: ProofOptions,
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> RescueRapsProver<H> {
    pub fn new(options: ProofOptions) -> Self {
        Self {
            options,
            _hasher: PhantomData,
        }
    }
    /// The parameter `seeds` is the set of seeds for the first hash chain.
    /// The parameter `permuted_seeds` is the set of seeds for the second hash chain.
    pub fn build_trace(
        &self,
        seeds: &[[BaseElement; 2]],
        permuted_seeds: &[[BaseElement; 2]],
        result: [[BaseElement; 2]; 2],
    ) -> RapTraceTable<BaseElement> {
        debug_assert_eq!(seeds.len(), permuted_seeds.len());
        // allocate memory to hold the trace table
        let trace_length = seeds.len() * CYCLE_LENGTH;
        let mut trace = RapTraceTable::new(2 * STATE_WIDTH, trace_length);
        const END_INCLUSIVE_RANGE: usize = NUM_HASH_ROUNDS - 1;

        trace.fill(
            |state| {
                // initialize original chain
                state[0] = seeds[0][0];
                state[1] = seeds[0][1];
                state[2] = BaseElement::ZERO;
                state[3] = BaseElement::ZERO;

                // initialize permuted chain
                state[4] = permuted_seeds[0][0];
                state[5] = permuted_seeds[0][1];
                state[6] = BaseElement::ZERO;
                state[7] = BaseElement::ZERO;
            },
            |step, state| {
                // execute the transition function for all steps
                //
                // for the first 14 steps in every cycle, compute a single round of
                // Rescue hash; for the remaining 2 rounds, carry over the values
                // in the first two registers of the two chains to the next step
                // and insert the additional seeds into the capacity registers
                match step % CYCLE_LENGTH {
                    0..=END_INCLUSIVE_RANGE => apply_rescue_round_parallel(state, step),
                    NUM_HASH_ROUNDS => {
                        let idx = step / CYCLE_LENGTH + 1;
                        // We don't have seeds for the final step once last hashing is done.
                        if idx < seeds.len() {
                            state[0] += seeds[idx][0];
                            state[1] += seeds[idx][1];

                            state[4] += permuted_seeds[idx][0];
                            state[5] += permuted_seeds[idx][1];
                        }
                    }
                    _ => {}
                };
            },
        );

        debug_assert_eq!(trace.get(0, trace_length - 1), result[0][0]);
        debug_assert_eq!(trace.get(1, trace_length - 1), result[0][1]);

        debug_assert_eq!(trace.get(4, trace_length - 1), result[1][0]);
        debug_assert_eq!(trace.get(5, trace_length - 1), result[1][1]);

        trace
    }
}

impl<H: ElementHasher> Prover for RescueRapsProver<H>
where
    H: ElementHasher<BaseField = BaseElement>,
{
    type BaseField = BaseElement;
    type Air = RescueRapsAir;
    type Trace = RapTraceTable<BaseElement>;
    type HashFn = H;
    type RandomCoin = DefaultRandomCoin<Self::HashFn>;

    fn get_pub_inputs(&self, trace: &Self::Trace) -> PublicInputs {
        let last_step = trace.length() - 1;
        PublicInputs {
            result: [
                [trace.get(0, last_step), trace.get(1, last_step)],
                [trace.get(4, last_step), trace.get(5, last_step)],
            ],
        }
    }

    fn options(&self) -> &ProofOptions {
        &self.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 super::Blake3_256;
use winterfell::{FieldExtension, ProofOptions};

#[test]
fn rescue_test_basic_proof_verification() {
    let rescue_eg = Box::new(super::RescueRapsExample::<Blake3_256>::new(
        128,
        build_options(false),
    ));
    crate::tests::test_basic_proof_verification(rescue_eg);
}

#[test]
fn rescue_test_basic_proof_verification_extension() {
    let rescue_eg = Box::new(super::RescueRapsExample::<Blake3_256>::new(
        128,
        build_options(true),
    ));
    crate::tests::test_basic_proof_verification(rescue_eg);
}

#[test]
fn rescue_test_basic_proof_verification_fail() {
    let rescue_eg = Box::new(super::RescueRapsExample::<Blake3_256>::new(
        128,
        build_options(false),
    ));
    crate::tests::test_basic_proof_verification_fail(rescue_eg);
}

fn build_options(use_extension_field: bool) -> ProofOptions {
    let extension = if use_extension_field {
        FieldExtension::Quadratic
    } else {
        FieldExtension::None
    };
    ProofOptions::new(28, 8, 0, extension, 4, 31)
}


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

use core::ops::Range;
use winterfell::{
    math::{FieldElement, StarkField},
    Trace, TraceTable,
};

pub mod rescue;

// CONSTRAINT EVALUATION HELPERS
// ================================================================================================

/// Returns zero only when a == b.
pub fn are_equal<E: FieldElement>(a: E, b: E) -> E {
    a - b
}

/// Returns zero only when a == zero.
pub fn is_zero<E: FieldElement>(a: E) -> E {
    a
}

/// Returns zero only when a = zero || a == one.
pub fn is_binary<E: FieldElement>(a: E) -> E {
    a * a - a
}

/// Return zero when a == one, and one when a == zero;
/// assumes that a is a binary value.
pub fn not<E: FieldElement>(a: E) -> E {
    E::ONE - a
}

// TRAIT TO SIMPLIFY CONSTRAINT AGGREGATION
// ================================================================================================

pub trait EvaluationResult<E> {
    fn agg_constraint(&mut self, index: usize, flag: E, value: E);
}

impl<E: FieldElement> EvaluationResult<E> for [E] {
    fn agg_constraint(&mut self, index: usize, flag: E, value: E) {
        self[index] += flag * value;
    }
}

impl<E: FieldElement> EvaluationResult<E> for Vec<E> {
    fn agg_constraint(&mut self, index: usize, flag: E, value: E) {
        self[index] += flag * value;
    }
}

// OTHER FUNCTIONS
// ================================================================================================

/// Prints out an execution trace.
pub fn print_trace<E: StarkField>(
    trace: &TraceTable<E>,
    multiples_of: usize,
    offset: usize,
    range: Range<usize>,
) {
    let trace_width = trace.width();

    let mut state = vec![E::ZERO; trace_width];
    for i in 0..trace.length() {
        if (i.wrapping_sub(offset)) % multiples_of != 0 {
            continue;
        }
        trace.read_row_into(i, &mut state);
        println!(
            "{}\t{:?}",
            i,
            state[range.clone()]
                .iter()
                .map(|v| v.as_int())
                .collect::<Vec<E::PositiveInteger>>()
        );
    }
}

pub fn print_trace_step<E: StarkField>(trace: &[Vec<E>], step: usize) {
    let trace_width = trace.len();
    let mut state = vec![E::ZERO; trace_width];
    for i in 0..trace_width {
        state[i] = trace[i][step];
    }
    println!(
        "{}\t{:?}",
        step,
        state
            .iter()
            .map(|v| v.as_int())
            .collect::<Vec<E::PositiveInteger>>()
    );
}


// 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::utils::{are_equal, EvaluationResult};
use core::slice;
use winterfell::{
    crypto::{Digest, Hasher},
    math::{fields::f128::BaseElement, FieldElement},
    ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
};

/// Function state is set to 6 field elements or 96 bytes; 4 elements are reserved for rate
/// and 2 elements are reserved for capacity.
pub const STATE_WIDTH: usize = 6;
pub const RATE_WIDTH: usize = 4;

/// Two elements (32-bytes) are returned as digest.
const DIGEST_SIZE: usize = 2;

/// The number of rounds is set to 7 to provide 128-bit security level with 40% security margin;
/// computed using algorithm 7 from <https://eprint.iacr.org/2020/1143.pdf>
/// security margin here differs from Rescue Prime specification which suggests 50% security
/// margin (and would require 8 rounds) primarily to make AIR a bit simpler.
pub const NUM_ROUNDS: usize = 7;

/// Minimum cycle length required to describe Rescue permutation.
pub const CYCLE_LENGTH: usize = 8;

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

pub struct Rescue128 {
    state: [BaseElement; STATE_WIDTH],
    idx: usize,
}

#[derive(Debug, Copy, Clone, Eq, PartialEq, Default)]
pub struct Hash([BaseElement; DIGEST_SIZE]);

// RESCUE128 IMPLEMENTATION
// ================================================================================================

impl Rescue128 {
    /// Returns a new hasher with the state initialized to all zeros.
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        Rescue128 {
            state: [BaseElement::ZERO; STATE_WIDTH],
            idx: 0,
        }
    }

    /// Absorbs data into the hasher state.
    pub fn update(&mut self, data: &[BaseElement]) {
        for &element in data {
            self.state[self.idx] += element;
            self.idx += 1;
            if self.idx % RATE_WIDTH == 0 {
                apply_permutation(&mut self.state);
                self.idx = 0;
            }
        }
    }

    /// Returns hash of the data absorbed into the hasher.
    pub fn finalize(mut self) -> Hash {
        if self.idx > 0 {
            // TODO: apply proper padding
            apply_permutation(&mut self.state);
        }
        Hash([self.state[0], self.state[1]])
    }

    /// Returns hash of the provided data.
    pub fn digest(data: &[BaseElement]) -> Hash {
        // initialize state to all zeros
        let mut state = [BaseElement::ZERO; STATE_WIDTH];

        let mut i = 0;
        for &element in data.iter() {
            state[i] += element;
            i += 1;
            if i % RATE_WIDTH == 0 {
                apply_permutation(&mut state);
                i = 0;
            }
        }

        if i > 0 {
            // TODO: apply proper padding
            apply_permutation(&mut state);
        }

        Hash([state[0], state[1]])
    }
}

// HASHER IMPLEMENTATION
// ================================================================================================

impl Hasher for Rescue128 {
    type Digest = Hash;

    const COLLISION_RESISTANCE: u32 = 64;

    fn hash(_bytes: &[u8]) -> Self::Digest {
        unimplemented!("not implemented")
    }

    fn merge(values: &[Self::Digest; 2]) -> Self::Digest {
        Self::digest(Hash::hashes_as_elements(values))
    }

    fn merge_with_int(_seed: Self::Digest, _value: u64) -> Self::Digest {
        unimplemented!("not implemented")
    }
}

// HASH IMPLEMENTATION
// ================================================================================================

impl Hash {
    pub fn new(v1: BaseElement, v2: BaseElement) -> Self {
        Hash([v1, v2])
    }

    #[allow(dead_code)]
    #[allow(clippy::wrong_self_convention)]
    pub fn to_bytes(&self) -> [u8; 32] {
        let mut bytes = [0; 32];
        bytes[..16].copy_from_slice(&self.0[0].to_bytes());
        bytes[16..].copy_from_slice(&self.0[1].to_bytes());
        bytes
    }

    #[allow(clippy::wrong_self_convention)]
    pub fn to_elements(&self) -> [BaseElement; DIGEST_SIZE] {
        self.0
    }

    pub fn hashes_as_elements(hashes: &[Hash]) -> &[BaseElement] {
        let p = hashes.as_ptr();
        let len = hashes.len() * DIGEST_SIZE;
        unsafe { slice::from_raw_parts(p as *const BaseElement, len) }
    }
}

impl Digest for Hash {
    fn as_bytes(&self) -> [u8; 32] {
        let bytes = BaseElement::elements_as_bytes(&self.0);
        let mut result = [0; 32];
        result[..bytes.len()].copy_from_slice(bytes);
        result
    }
}

impl Serializable for Hash {
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        target.write(self.0[0]);
        target.write(self.0[1]);
    }
}

impl Deserializable for Hash {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let v1 = BaseElement::read_from(source)?;
        let v2 = BaseElement::read_from(source)?;
        Ok(Self([v1, v2]))
    }
}

// RESCUE PERMUTATION
// ================================================================================================

/// Applies Rescue-XLIX permutation to the provided state.
pub fn apply_permutation(state: &mut [BaseElement; STATE_WIDTH]) {
    // apply round function 7 times; this provides 128-bit security with 40% security margin
    for i in 0..NUM_ROUNDS {
        apply_round(state, i);
    }
}

/// Rescue-XLIX round function;
/// implementation based on algorithm 3 from <https://eprint.iacr.org/2020/1143.pdf>
#[inline(always)]
pub fn apply_round(state: &mut [BaseElement], step: usize) {
    // determine which round constants to use
    let ark = ARK[step % CYCLE_LENGTH];

    // apply first half of Rescue round
    apply_sbox(state);
    apply_mds(state);
    for i in 0..STATE_WIDTH {
        state[i] += ark[i];
    }

    // apply second half of Rescue round
    apply_inv_sbox(state);
    apply_mds(state);
    for i in 0..STATE_WIDTH {
        state[i] += ark[STATE_WIDTH + i];
    }
}

// CONSTRAINTS
// ================================================================================================

/// when flag = 1, enforces constraints for a single round of Rescue hash functions
pub fn enforce_round<E: FieldElement + From<BaseElement>>(
    result: &mut [E],
    current: &[E],
    next: &[E],
    ark: &[E],
    flag: E,
) {
    // compute the state that should result from applying the first half of Rescue round
    // to the current state of the computation
    let mut step1 = [E::ZERO; STATE_WIDTH];
    step1.copy_from_slice(current);
    apply_sbox(&mut step1);
    apply_mds(&mut step1);
    for i in 0..STATE_WIDTH {
        step1[i] += ark[i];
    }

    // compute the state that should result from applying the inverse for the second
    // half for Rescue round to the next step of the computation
    let mut step2 = [E::ZERO; STATE_WIDTH];
    step2.copy_from_slice(next);
    for i in 0..STATE_WIDTH {
        step2[i] -= ark[STATE_WIDTH + i];
    }
    apply_inv_mds(&mut step2);
    apply_sbox(&mut step2);

    // make sure that the results are equal
    for i in 0..STATE_WIDTH {
        result.agg_constraint(i, flag, are_equal(step2[i], step1[i]));
    }
}

// ROUND CONSTANTS
// ================================================================================================

/// Returns Rescue round constants arranged in column-major form.
pub fn get_round_constants() -> Vec<Vec<BaseElement>> {
    let mut constants = Vec::new();
    for _ in 0..(STATE_WIDTH * 2) {
        constants.push(vec![BaseElement::ZERO; CYCLE_LENGTH]);
    }

    #[allow(clippy::needless_range_loop)]
    for i in 0..CYCLE_LENGTH {
        for j in 0..(STATE_WIDTH * 2) {
            constants[j][i] = ARK[i][j];
        }
    }

    constants
}

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

#[inline(always)]
#[allow(clippy::needless_range_loop)]
fn apply_sbox<E: FieldElement>(state: &mut [E]) {
    for i in 0..STATE_WIDTH {
        state[i] = state[i].exp(ALPHA.into());
    }
}

#[inline(always)]
#[allow(clippy::needless_range_loop)]
fn apply_inv_sbox(state: &mut [BaseElement]) {
    for i in 0..STATE_WIDTH {
        state[i] = state[i].exp(INV_ALPHA);
    }
}

#[inline(always)]
#[allow(clippy::needless_range_loop)]
fn apply_mds<E: FieldElement + From<BaseElement>>(state: &mut [E]) {
    let mut result = [E::ZERO; STATE_WIDTH];
    let mut temp = [E::ZERO; STATE_WIDTH];
    for i in 0..STATE_WIDTH {
        for j in 0..STATE_WIDTH {
            temp[j] = E::from(MDS[i * STATE_WIDTH + j]) * state[j];
        }

        for j in 0..STATE_WIDTH {
            result[i] += temp[j];
        }
    }
    state.copy_from_slice(&result);
}

#[inline(always)]
#[allow(clippy::needless_range_loop)]
fn apply_inv_mds<E: FieldElement + From<BaseElement>>(state: &mut [E]) {
    let mut result = [E::ZERO; STATE_WIDTH];
    let mut temp = [E::ZERO; STATE_WIDTH];
    for i in 0..STATE_WIDTH {
        for j in 0..STATE_WIDTH {
            temp[j] = E::from(INV_MDS[i * STATE_WIDTH + j]) * state[j];
        }

        for j in 0..STATE_WIDTH {
            result[i] += temp[j];
        }
    }
    state.copy_from_slice(&result);
}

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

/// S-Box and Inverse S-Box powers;
/// computed using algorithm 6 from <https://eprint.iacr.org/2020/1143.pdf>
const ALPHA: u32 = 5;
const INV_ALPHA: u128 = 272225893536750770770699646362995969229;

/// Rescue MDS matrix
/// Computed using algorithm 4 from <https://eprint.iacr.org/2020/1143.pdf>
const MDS: [BaseElement; STATE_WIDTH * STATE_WIDTH] = [
    BaseElement::new(340282366920938463463374557953730612630),
    BaseElement::new(21493836),
    BaseElement::new(340282366920938463463374557953736934518),
    BaseElement::new(914760),
    BaseElement::new(340282366920938463463374557953744928504),
    BaseElement::new(364),
    BaseElement::new(340282366920938463463374557948521959389),
    BaseElement::new(7809407397),
    BaseElement::new(340282366920938463463374557950844620457),
    BaseElement::new(324945621),
    BaseElement::new(340282366920938463463374557953733852285),
    BaseElement::new(99463),
    BaseElement::new(340282366920938463463374556526559624596),
    BaseElement::new(2132618407920),
    BaseElement::new(340282366920938463463374557163162978137),
    BaseElement::new(88084432800),
    BaseElement::new(340282366920938463463374557950784345879),
    BaseElement::new(25095280),
    BaseElement::new(340282366920938463463374197863906102577),
    BaseElement::new(537966647357139),
    BaseElement::new(340282366920938463463374358646073999137),
    BaseElement::new(22165576349400),
    BaseElement::new(340282366920938463463374557212857010097),
    BaseElement::new(6174066262),
    BaseElement::new(340282366920938463463285966851139685903),
    BaseElement::new(132344277849702072),
    BaseElement::new(340282366920938463463325536573199985698),
    BaseElement::new(5448481182864720),
    BaseElement::new(340282366920938463463374376171390478291),
    BaseElement::new(1506472167928),
    BaseElement::new(340282366920938463441758328918057706841),
    BaseElement::new(32291274613403616174),
    BaseElement::new(340282366920938463451414421516665416977),
    BaseElement::new(1329039099788841441),
    BaseElement::new(340282366920938463463330243139804660633),
    BaseElement::new(366573514642546),
];

const INV_MDS: [BaseElement; STATE_WIDTH * STATE_WIDTH] = [
    BaseElement::new(133202720344903784697302507504318451498),
    BaseElement::new(9109562341901685402869515497167051415),
    BaseElement::new(187114562320006661061623258692072377978),
    BaseElement::new(217977550980311337650875125151512141987),
    BaseElement::new(274535269264332978809051716514493438195),
    BaseElement::new(198907435511358942768401550501671423539),
    BaseElement::new(107211340690419935719675873160429442610),
    BaseElement::new(93035459208639798096019355873148696692),
    BaseElement::new(34612840942819361370119536876515785819),
    BaseElement::new(28124271099756519590702162721340502811),
    BaseElement::new(220883180661145883341796932300840696133),
    BaseElement::new(196697641239095428808435254975214799010),
    BaseElement::new(48198755643822249649260269442324041679),
    BaseElement::new(64419499747985404280993270855996080557),
    BaseElement::new(280207800449835933431237404716657540948),
    BaseElement::new(61755931245950637038951462642746253929),
    BaseElement::new(206737575380416523686108693210925354496),
    BaseElement::new(19245171373866178840198015038840651466),
    BaseElement::new(133290479635000282395391929030461304726),
    BaseElement::new(256035933367497105928763702353648605000),
    BaseElement::new(97077987470620839632334052346344687963),
    BaseElement::new(144638736603246051821344039641662500876),
    BaseElement::new(323753713558453221824969500168839490204),
    BaseElement::new(66050250127997888787320450320278295843),
    BaseElement::new(107416947271017171483976049725783774207),
    BaseElement::new(29799978553141132526384006297614595309),
    BaseElement::new(112991183841517485419461727429810868869),
    BaseElement::new(27096959906733835564333321118624482460),
    BaseElement::new(197262955506413467422574209301409294301),
    BaseElement::new(205996708763053834510019802034246907929),
    BaseElement::new(114827794598835201662537916675749328586),
    BaseElement::new(22232541983454090535685600849896663137),
    BaseElement::new(84718265936029339288536427868493390350),
    BaseElement::new(176534200716138685131361645691447579493),
    BaseElement::new(304590074876810806644622682832255680729),
    BaseElement::new(317944222651547267127379399943392242317),
];

/// Rescue round constants;
/// computed using algorithm 5 from <https://eprint.iacr.org/2020/1143.pdf>
pub const ARK: [[BaseElement; STATE_WIDTH * 2]; CYCLE_LENGTH] = [
    [
        BaseElement::new(232350694689151131917165570858777669544),
        BaseElement::new(297138716840883070166239111380460167036),
        BaseElement::new(262280230220923724082396709497064092149),
        BaseElement::new(172158049344191113832187131208632037738),
        BaseElement::new(49064466045797039562408393043269857959),
        BaseElement::new(310779117230843293557874990285120450495),
        BaseElement::new(256706820970445617734149759518940865107),
        BaseElement::new(79123538858040670180278455836284339197),
        BaseElement::new(78750303544367952484014721485273250812),
        BaseElement::new(288861383492149579433903883762711410179),
        BaseElement::new(59801749333456280387477464033868461625),
        BaseElement::new(21443300235508431203706748477819269958),
    ],
    [
        BaseElement::new(58568963110264836729315799795504150465),
        BaseElement::new(330748576252425315826992430477036516321),
        BaseElement::new(186265990460580587588657915966473647991),
        BaseElement::new(33474186560709631768594728335471560699),
        BaseElement::new(158848462530608412921046130349797355353),
        BaseElement::new(103951280788776493556470655637893338265),
        BaseElement::new(143328281743837680325887693977200434046),
        BaseElement::new(84141533915622931968833899936597847300),
        BaseElement::new(8289043147167319381038668861607412243),
        BaseElement::new(182690551456641207603161012621368395791),
        BaseElement::new(189966993584382842241685332212477020587),
        BaseElement::new(32137923394454105763485467845755642950),
    ],
    [
        BaseElement::new(37831789571282423629213813309051107559),
        BaseElement::new(128553631204082467137622394929811125529),
        BaseElement::new(267986778741944677472811189878493395927),
        BaseElement::new(16604948458564067211433039503683613987),
        BaseElement::new(336102510949899388907937615764984494068),
        BaseElement::new(269515689098362827313089599343791905108),
        BaseElement::new(299424679105391259942771484229152481303),
        BaseElement::new(204910193356347483970850685012209050540),
        BaseElement::new(297547986861132400067173315704469727918),
        BaseElement::new(90994669428470088728996184833134573519),
        BaseElement::new(194832530917116381832912394976136685925),
        BaseElement::new(3544879195102182108390682435201981399),
    ],
    [
        BaseElement::new(339480205126523778084089852053600037139),
        BaseElement::new(7584482258985997923597941079175892345),
        BaseElement::new(293411952222390873312400094181647328549),
        BaseElement::new(199529004542042321671242096609546451065),
        BaseElement::new(67129123347758775813781826519244753478),
        BaseElement::new(262358775581253675478636059962684988488),
        BaseElement::new(214578730175648891816936630380713062555),
        BaseElement::new(298888476681892954783673663609236117055),
        BaseElement::new(28713802418311531156758766332916445632),
        BaseElement::new(1440134829402109711440873134882900954),
        BaseElement::new(136568912729847804743104940208565395935),
        BaseElement::new(282333114631262903665175684297593586626),
    ],
    [
        BaseElement::new(179980515973143677823617972256218090691),
        BaseElement::new(262324617228293661450608983002445445851),
        BaseElement::new(101457408539557988072857167265007764003),
        BaseElement::new(135015365700146217343913438445165565670),
        BaseElement::new(160037359781136723784361845515476884821),
        BaseElement::new(182530253870899012049936279038476084254),
        BaseElement::new(135879876810809726132885131537021449499),
        BaseElement::new(232021530889024386996643355214152586646),
        BaseElement::new(145764181560102807472161589832442506602),
        BaseElement::new(30096323905520593555387863391076216460),
        BaseElement::new(26964230850883304384940372063347292502),
        BaseElement::new(248723932438838238159920468579438468564),
    ],
    [
        BaseElement::new(294269904099379916907622037481357861347),
        BaseElement::new(68547751515194812125080398554316505804),
        BaseElement::new(206967528806115588933607920597265054243),
        BaseElement::new(218563991130423186053843420486943196637),
        BaseElement::new(271753381570791699387473121354016967661),
        BaseElement::new(280821616954361601859332610476339898658),
        BaseElement::new(10004341245328361103806488533574675264),
        BaseElement::new(102737972201824925757345477497905200949),
        BaseElement::new(181579715086871199454198713448655357907),
        BaseElement::new(334443686013848360201749831728546200670),
        BaseElement::new(43930702221243327593116820380585481596),
        BaseElement::new(16744004758332429127464852702179311517),
    ],
    [
        BaseElement::new(310201738135125726809998762242791360596),
        BaseElement::new(155126893730515639579436939964032992002),
        BaseElement::new(61238650483248463229462616021804212788),
        BaseElement::new(6693212157784826508674787451860949238),
        BaseElement::new(197651057967963974372308220503477603713),
        BaseElement::new(174221476673212934077040088950046690415),
        BaseElement::new(287511813733819668564695051918836002922),
        BaseElement::new(304531189544765525159398110881793396421),
        BaseElement::new(276777415462914862553995344360435589651),
        BaseElement::new(241036817921529641113885285343669990717),
        BaseElement::new(320958231309550951576801366383624382828),
        BaseElement::new(242260690344880997681123448650535822378),
    ],
    [
        BaseElement::new(201589105262974747061391276271612166799),
        BaseElement::new(21009766855942890883171267876432289297),
        BaseElement::new(303226336248222109074995022589884483065),
        BaseElement::new(105515432862530091605210357969101266504),
        BaseElement::new(235097661610089805414814372959229370626),
        BaseElement::new(210361497167001742816223425802317150493),
        BaseElement::new(218546747003262668455051521918398855294),
        BaseElement::new(280724473534270362895764829061545243190),
        BaseElement::new(179926408118748249708833901850481685351),
        BaseElement::new(168859451670725335987760025077648496937),
        BaseElement::new(127174659756870191527945451601624140498),
        BaseElement::new(290826558340641225374953827677533570165),
    ],
];


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

pub mod exempt;
pub mod regular;


// 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::{BaseElement, FieldElement, ProofOptions, ALPHA, FORTY_TWO};
use winterfell::{
    math::ToElements, Air, AirContext, Assertion, EvaluationFrame, TraceInfo,
    TransitionConstraintDegree,
};

// PUBLIC INPUTS
// ================================================================================================

#[derive(Clone)]
pub struct VdfInputs {
    pub seed: BaseElement,
    pub result: BaseElement,
}

impl ToElements<BaseElement> for VdfInputs {
    fn to_elements(&self) -> Vec<BaseElement> {
        vec![self.seed, self.result]
    }
}

// VDF AIR
// ================================================================================================

pub struct VdfAir {
    context: AirContext<BaseElement>,
    seed: BaseElement,
    result: BaseElement,
}

impl Air for VdfAir {
    type BaseField = BaseElement;
    type PublicInputs = VdfInputs;

    fn new(trace_info: TraceInfo, pub_inputs: VdfInputs, options: ProofOptions) -> Self {
        let degrees = vec![TransitionConstraintDegree::new(3)];
        // make sure the last two rows are excluded from transition constraints as we populate
        // values in the last row with garbage
        let context =
            AirContext::new(trace_info, degrees, 2, options).set_num_transition_exemptions(2);
        Self {
            context,
            seed: pub_inputs.seed,
            result: pub_inputs.result,
        }
    }

    fn evaluate_transition<E: FieldElement<BaseField = Self::BaseField>>(
        &self,
        frame: &EvaluationFrame<E>,
        _periodic_values: &[E],
        result: &mut [E],
    ) {
        let current_state = frame.current()[0];
        let next_state = frame.next()[0];

        result[0] = current_state - (next_state.exp(ALPHA.into()) + FORTY_TWO.into());
    }

    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
        // the second boundary constraint is applied to the second to last step
        // as the last step will contain garbage
        let second_to_last_step = self.trace_length() - 2;
        vec![
            Assertion::single(0, 0, self.seed),
            Assertion::single(0, second_to_last_step, self.result),
        ]
    }

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


// 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::{Blake3_192, Blake3_256, Example, ExampleOptions, HashFunction, Sha3_256};
use core::marker::PhantomData;
use log::debug;
use std::time::Instant;
use winterfell::{
    crypto::{DefaultRandomCoin, ElementHasher},
    math::{fields::f128::BaseElement, FieldElement},
    ProofOptions, Prover, StarkProof, Trace, TraceTable, VerifierError,
};

mod air;
use air::{VdfAir, VdfInputs};

mod prover;
use prover::VdfProver;

#[cfg(test)]
mod tests;

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

const ALPHA: u64 = 3;
const INV_ALPHA: u128 = 226854911280625642308916371969163307691;
const FORTY_TWO: BaseElement = BaseElement::new(42);

// VDF EXAMPLE
// ================================================================================================

pub fn get_example(options: &ExampleOptions, num_steps: usize) -> Result<Box<dyn Example>, String> {
    let (options, hash_fn) = options.to_proof_options(85, 2);

    match hash_fn {
        HashFunction::Blake3_192 => Ok(Box::new(VdfExample::<Blake3_192>::new(num_steps, options))),
        HashFunction::Blake3_256 => Ok(Box::new(VdfExample::<Blake3_256>::new(num_steps, options))),
        HashFunction::Sha3_256 => Ok(Box::new(VdfExample::<Sha3_256>::new(num_steps, options))),
        _ => Err("The specified hash function cannot be used with this example.".to_string()),
    }
}

pub struct VdfExample<H: ElementHasher> {
    options: ProofOptions,
    num_steps: usize,
    seed: BaseElement,
    result: BaseElement,
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> VdfExample<H> {
    pub fn new(num_steps: usize, options: ProofOptions) -> Self {
        assert!(
            (num_steps + 1).is_power_of_two(),
            "number of steps must be one less than a power of 2"
        );

        // run the VDF function
        let now = Instant::now();
        let seed = BaseElement::new(123);
        let result = execute_vdf(seed, num_steps);
        debug!(
            "Executed the VDF function for {} steps in {} ms",
            num_steps,
            now.elapsed().as_millis()
        );

        Self {
            options,
            num_steps,
            seed,
            result,
            _hasher: PhantomData,
        }
    }
}

// EXAMPLE IMPLEMENTATION
// ================================================================================================

impl<H: ElementHasher> Example for VdfExample<H>
where
    H: ElementHasher<BaseField = BaseElement>,
{
    fn prove(&self) -> StarkProof {
        debug!(
            "Generating proof for executing a VDF function for {} steps\n\
            ---------------------",
            self.num_steps
        );

        // create a prover
        let prover = VdfProver::<H>::new(self.options.clone());

        // generate execution trace
        let now = Instant::now();
        let trace = VdfProver::<H>::build_trace(self.seed, self.num_steps + 1);

        let trace_width = trace.width();
        let trace_length = trace.length();
        debug!(
            "Generated execution trace of {} registers and 2^{} steps in {} ms",
            trace_width,
            trace_length.ilog2(),
            now.elapsed().as_millis()
        );

        // generate the proof
        prover.prove(trace).unwrap()
    }

    fn verify(&self, proof: StarkProof) -> Result<(), VerifierError> {
        let pub_inputs = VdfInputs {
            seed: self.seed,
            result: self.result,
        };
        winterfell::verify::<VdfAir, H, DefaultRandomCoin<H>>(proof, pub_inputs)
    }

    fn verify_with_wrong_inputs(&self, proof: StarkProof) -> Result<(), VerifierError> {
        let pub_inputs = VdfInputs {
            seed: self.seed,
            result: self.result + BaseElement::ONE,
        };
        winterfell::verify::<VdfAir, H, DefaultRandomCoin<H>>(proof, pub_inputs)
    }
}

// VDF FUNCTION
// ================================================================================================

fn execute_vdf(seed: BaseElement, n: usize) -> BaseElement {
    let mut state = seed;
    for _ in 0..(n - 1) {
        state = (state - FORTY_TWO).exp(INV_ALPHA);
    }
    state
}


// 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::{
    BaseElement, DefaultRandomCoin, ElementHasher, FieldElement, PhantomData, ProofOptions, Prover,
    Trace, TraceTable, VdfAir, VdfInputs, FORTY_TWO, INV_ALPHA,
};

// VDF PROVER
// ================================================================================================

pub struct VdfProver<H: ElementHasher> {
    options: ProofOptions,
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> VdfProver<H> {
    pub fn new(options: ProofOptions) -> Self {
        Self {
            options,
            _hasher: PhantomData,
        }
    }

    pub fn build_trace(seed: BaseElement, n: usize) -> TraceTable<BaseElement> {
        let mut trace = Vec::with_capacity(n);
        let mut state = seed;

        trace.push(state);
        for _ in 0..(n - 2) {
            state = (state - FORTY_TWO).exp(INV_ALPHA);
            trace.push(state);
        }

        // put garbage value into the last step
        trace.push(BaseElement::new(123));

        TraceTable::init(vec![trace])
    }
}

impl<H: ElementHasher> Prover for VdfProver<H>
where
    H: ElementHasher<BaseField = BaseElement>,
{
    type BaseField = BaseElement;
    type Air = VdfAir;
    type Trace = TraceTable<BaseElement>;
    type HashFn = H;
    type RandomCoin = DefaultRandomCoin<Self::HashFn>;

    fn get_pub_inputs(&self, trace: &Self::Trace) -> VdfInputs {
        // the result is read from the second to last step because the last last step contains
        // garbage
        let second_to_last_step = trace.length() - 2;
        VdfInputs {
            seed: trace.get(0, 0),
            result: trace.get(0, second_to_last_step),
        }
    }

    fn options(&self) -> &ProofOptions {
        &self.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 super::Blake3_256;
use winterfell::{FieldExtension, ProofOptions};

#[test]
fn vdf_test_basic_proof_verification() {
    let fib = Box::new(super::VdfExample::<Blake3_256>::new(
        127,
        build_options(false),
    ));
    crate::tests::test_basic_proof_verification(fib);
}

#[test]
fn vdf_test_basic_proof_verification_extension() {
    let fib = Box::new(super::VdfExample::<Blake3_256>::new(
        127,
        build_options(true),
    ));
    crate::tests::test_basic_proof_verification(fib);
}

#[test]
fn vdf_test_basic_proof_verification_fail() {
    let fib = Box::new(super::VdfExample::<Blake3_256>::new(
        127,
        build_options(false),
    ));
    crate::tests::test_basic_proof_verification_fail(fib);
}

fn build_options(use_extension_field: bool) -> ProofOptions {
    let extension = if use_extension_field {
        FieldExtension::Quadratic
    } else {
        FieldExtension::None
    };
    ProofOptions::new(85, 2, 0, extension, 4, 31)
}


// 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::{BaseElement, FieldElement, ProofOptions, ALPHA, FORTY_TWO};
use winterfell::{
    math::ToElements, Air, AirContext, Assertion, EvaluationFrame, TraceInfo,
    TransitionConstraintDegree,
};

// PUBLIC INPUTS
// ================================================================================================

#[derive(Clone)]
pub struct VdfInputs {
    pub seed: BaseElement,
    pub result: BaseElement,
}

impl ToElements<BaseElement> for VdfInputs {
    fn to_elements(&self) -> Vec<BaseElement> {
        vec![self.seed, self.result]
    }
}

// VDF AIR
// ================================================================================================

pub struct VdfAir {
    context: AirContext<BaseElement>,
    seed: BaseElement,
    result: BaseElement,
}

impl Air for VdfAir {
    type BaseField = BaseElement;
    type PublicInputs = VdfInputs;

    fn new(trace_info: TraceInfo, pub_inputs: VdfInputs, options: ProofOptions) -> Self {
        let degrees = vec![TransitionConstraintDegree::new(3)];
        Self {
            context: AirContext::new(trace_info, degrees, 2, options),
            seed: pub_inputs.seed,
            result: pub_inputs.result,
        }
    }

    fn evaluate_transition<E: FieldElement<BaseField = Self::BaseField>>(
        &self,
        frame: &EvaluationFrame<E>,
        _periodic_values: &[E],
        result: &mut [E],
    ) {
        let current_state = frame.current()[0];
        let next_state = frame.next()[0];

        result[0] = current_state - (next_state.exp(ALPHA.into()) + FORTY_TWO.into());
    }

    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
        let last_step = self.trace_length() - 1;
        vec![
            Assertion::single(0, 0, self.seed),
            Assertion::single(0, last_step, self.result),
        ]
    }

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


// 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::{Blake3_192, Blake3_256, Example, ExampleOptions, HashFunction, Sha3_256};
use core::marker::PhantomData;
use log::debug;
use std::time::Instant;
use winterfell::{
    crypto::{DefaultRandomCoin, ElementHasher},
    math::{fields::f128::BaseElement, FieldElement},
    ProofOptions, Prover, StarkProof, Trace, TraceTable, VerifierError,
};

mod air;
use air::{VdfAir, VdfInputs};

mod prover;
use prover::VdfProver;

#[cfg(test)]
mod tests;

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

const ALPHA: u64 = 3;
const INV_ALPHA: u128 = 226854911280625642308916371969163307691;
const FORTY_TWO: BaseElement = BaseElement::new(42);

// VDF EXAMPLE
// ================================================================================================

pub fn get_example(options: &ExampleOptions, num_steps: usize) -> Result<Box<dyn Example>, String> {
    let (options, hash_fn) = options.to_proof_options(85, 2);

    match hash_fn {
        HashFunction::Blake3_192 => Ok(Box::new(VdfExample::<Blake3_192>::new(num_steps, options))),
        HashFunction::Blake3_256 => Ok(Box::new(VdfExample::<Blake3_256>::new(num_steps, options))),
        HashFunction::Sha3_256 => Ok(Box::new(VdfExample::<Sha3_256>::new(num_steps, options))),
        _ => Err("The specified hash function cannot be used with this example.".to_string()),
    }
}

pub struct VdfExample<H: ElementHasher> {
    options: ProofOptions,
    num_steps: usize,
    seed: BaseElement,
    result: BaseElement,
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> VdfExample<H> {
    pub fn new(num_steps: usize, options: ProofOptions) -> Self {
        assert!(
            num_steps.is_power_of_two(),
            "number of steps must be a power of 2"
        );

        // run the VDF function
        let now = Instant::now();
        let seed = BaseElement::new(123);
        let result = execute_vdf(seed, num_steps);
        debug!(
            "Executed the VDF function for {} steps in {} ms",
            num_steps,
            now.elapsed().as_millis()
        );

        Self {
            options,
            num_steps,
            seed,
            result,
            _hasher: PhantomData,
        }
    }
}

// EXAMPLE IMPLEMENTATION
// ================================================================================================

impl<H: ElementHasher> Example for VdfExample<H>
where
    H: ElementHasher<BaseField = BaseElement>,
{
    fn prove(&self) -> StarkProof {
        debug!(
            "Generating proof for executing a VDF function for {} steps\n\
            ---------------------",
            self.num_steps
        );

        // create a prover
        let prover = VdfProver::<H>::new(self.options.clone());

        // generate execution trace
        let now = Instant::now();
        let trace = VdfProver::<H>::build_trace(self.seed, self.num_steps);

        let trace_width = trace.width();
        let trace_length = trace.length();
        debug!(
            "Generated execution trace of {} registers and 2^{} steps in {} ms",
            trace_width,
            trace_length.ilog2(),
            now.elapsed().as_millis()
        );

        // generate the proof
        prover.prove(trace).unwrap()
    }

    fn verify(&self, proof: StarkProof) -> Result<(), VerifierError> {
        let pub_inputs = VdfInputs {
            seed: self.seed,
            result: self.result,
        };
        winterfell::verify::<VdfAir, H, DefaultRandomCoin<H>>(proof, pub_inputs)
    }

    fn verify_with_wrong_inputs(&self, proof: StarkProof) -> Result<(), VerifierError> {
        let pub_inputs = VdfInputs {
            seed: self.seed,
            result: self.result + BaseElement::ONE,
        };
        winterfell::verify::<VdfAir, H, DefaultRandomCoin<H>>(proof, pub_inputs)
    }
}

// VDF FUNCTION
// ================================================================================================

fn execute_vdf(seed: BaseElement, n: usize) -> BaseElement {
    let mut state = seed;
    for _ in 0..(n - 1) {
        state = (state - FORTY_TWO).exp(INV_ALPHA);
    }
    state
}


// 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::{
    BaseElement, DefaultRandomCoin, ElementHasher, FieldElement, PhantomData, ProofOptions, Prover,
    Trace, TraceTable, VdfAir, VdfInputs, FORTY_TWO, INV_ALPHA,
};

// VDF PROVER
// ================================================================================================

pub struct VdfProver<H: ElementHasher> {
    options: ProofOptions,
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> VdfProver<H> {
    pub fn new(options: ProofOptions) -> Self {
        Self {
            options,
            _hasher: PhantomData,
        }
    }

    pub fn build_trace(seed: BaseElement, n: usize) -> TraceTable<BaseElement> {
        let mut trace = Vec::with_capacity(n);
        let mut state = seed;

        trace.push(state);
        for _ in 0..(n - 1) {
            state = (state - FORTY_TWO).exp(INV_ALPHA);
            trace.push(state);
        }

        TraceTable::init(vec![trace])
    }
}

impl<H: ElementHasher> Prover for VdfProver<H>
where
    H: ElementHasher<BaseField = BaseElement>,
{
    type BaseField = BaseElement;
    type Air = VdfAir;
    type Trace = TraceTable<BaseElement>;
    type HashFn = H;
    type RandomCoin = DefaultRandomCoin<Self::HashFn>;

    fn get_pub_inputs(&self, trace: &Self::Trace) -> VdfInputs {
        let last_step = trace.length() - 1;
        VdfInputs {
            seed: trace.get(0, 0),
            result: trace.get(0, last_step),
        }
    }

    fn options(&self) -> &ProofOptions {
        &self.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 super::Blake3_256;
use winterfell::{FieldExtension, ProofOptions};

#[test]
fn vdf_test_basic_proof_verification() {
    let fib = Box::new(super::VdfExample::<Blake3_256>::new(
        128,
        build_options(false),
    ));
    crate::tests::test_basic_proof_verification(fib);
}

#[test]
fn vdf_test_basic_proof_verification_extension() {
    let fib = Box::new(super::VdfExample::<Blake3_256>::new(
        128,
        build_options(true),
    ));
    crate::tests::test_basic_proof_verification(fib);
}

#[test]
fn vdf_test_basic_proof_verification_fail() {
    let fib = Box::new(super::VdfExample::<Blake3_256>::new(
        128,
        build_options(false),
    ));
    crate::tests::test_basic_proof_verification_fail(fib);
}

fn build_options(use_extension_field: bool) -> ProofOptions {
    let extension = if use_extension_field {
        FieldExtension::Quadratic
    } else {
        FieldExtension::None
    };
    ProofOptions::new(85, 2, 0, extension, 4, 31)
}


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

use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use math::{fields::f128::BaseElement, get_power_series, polynom, StarkField};
use rand_utils::{rand_value, rand_vector};
use utils::group_vector_elements;
use winter_fri::folding;

static BATCH_SIZES: [usize; 3] = [65536, 131072, 262144];

pub fn interpolate_batch(c: &mut Criterion) {
    let mut group = c.benchmark_group("interpolate batch");

    for &size in &BATCH_SIZES {
        let (xs, ys) = build_coordinate_batches(size);
        group.bench_function(BenchmarkId::new("generic", size), |b| {
            b.iter(|| polynom::interpolate_batch(&xs, &ys))
        });
    }
}

pub fn apply_drp(c: &mut Criterion) {
    let mut group = c.benchmark_group("drp");

    for &size in &BATCH_SIZES {
        let (_, ys) = build_coordinate_batches(size);
        let alpha: BaseElement = rand_value();
        group.bench_function(BenchmarkId::new("base field", size), |b| {
            b.iter(|| folding::apply_drp(&ys, BaseElement::GENERATOR, alpha))
        });
    }
}

criterion_group!(quartic_group, interpolate_batch, apply_drp);
criterion_main!(quartic_group);

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

fn build_coordinate_batches(batch_size: usize) -> (Vec<[BaseElement; 4]>, Vec<[BaseElement; 4]>) {
    let r = BaseElement::get_root_of_unity(batch_size.ilog2());
    let xs = group_vector_elements(get_power_series(r, batch_size));
    let ys = group_vector_elements(rand_vector::<BaseElement>(batch_size));
    (xs, ys)
}


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

use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
use crypto::{hashers::Blake3_256, DefaultRandomCoin};
use math::{fft, fields::f128::BaseElement, FieldElement};
use rand_utils::rand_vector;
use std::time::Duration;
use winter_fri::{DefaultProverChannel, FriOptions, FriProver};

static BATCH_SIZES: [usize; 3] = [65536, 131072, 262144];
static BLOWUP_FACTOR: usize = 8;

pub fn build_layers(c: &mut Criterion) {
    let mut fri_group = c.benchmark_group("FRI prover");
    fri_group.sample_size(10);
    fri_group.measurement_time(Duration::from_secs(10));

    let options = FriOptions::new(BLOWUP_FACTOR, 4, 255);

    for &domain_size in &BATCH_SIZES {
        let evaluations = build_evaluations(domain_size);

        fri_group.bench_with_input(
            BenchmarkId::new("build_layers", domain_size),
            &evaluations,
            |b, e| {
                let mut prover = FriProver::new(options.clone());
                b.iter_batched(
                    || e.clone(),
                    |evaluations| {
                        let mut channel = DefaultProverChannel::<
                            BaseElement,
                            Blake3_256<BaseElement>,
                            DefaultRandomCoin<Blake3_256<BaseElement>>,
                        >::new(domain_size, 32);
                        prover.build_layers(&mut channel, evaluations);
                        prover.reset();
                    },
                    BatchSize::LargeInput,
                );
            },
        );
    }
}

criterion_group!(fri_prover_group, build_layers);
criterion_main!(fri_prover_group);

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

fn build_evaluations(domain_size: usize) -> Vec<BaseElement> {
    let mut p: Vec<BaseElement> = rand_vector(domain_size / BLOWUP_FACTOR);
    p.resize(domain_size, BaseElement::ZERO);
    let twiddles = fft::get_twiddles::<BaseElement>(domain_size);
    fft::evaluate_poly(&mut p, &twiddles);
    p
}


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

use core::fmt;

use crypto::RandomCoinError;

// VERIFIER ERROR
// ================================================================================================

/// Defines errors which can occur during FRI proof verification.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum VerifierError {
    /// Attempt to draw a random value from a public coin failed.
    RandomCoinError(RandomCoinError),
    /// Folding factor specified for the protocol is not supported. Currently, supported folding
    /// factors are: 4, 8, and 16.
    UnsupportedFoldingFactor(usize),
    /// Number of query positions does not match the number of provided evaluations.
    NumPositionEvaluationMismatch(usize, usize),
    /// Evaluations at queried positions did not match layer commitment made by the prover.
    LayerCommitmentMismatch,
    /// Degree-respecting projection was not performed correctly at one of the layers.
    InvalidLayerFolding(usize),
    /// FRI remainder did not match the commitment.
    RemainderCommitmentMismatch,
    /// Degree-respecting projection was not performed correctly at the last layer.
    InvalidRemainderFolding,
    /// FRI remainder expected degree is greater than number of remainder values.
    RemainderDegreeNotValid,
    /// FRI remainder degree is greater than the polynomial degree expected for the last layer.
    RemainderDegreeMismatch(usize),
    /// Polynomial degree at one of the FRI layers could not be divided evenly by the folding factor.
    DegreeTruncation(usize, usize, usize),
}

impl fmt::Display for VerifierError {
    #[rustfmt::skip]
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::RandomCoinError(err) => {
                write!(f, "failed to draw a random value from the public coin: {err}")
            }
            Self::UnsupportedFoldingFactor(value) => {
                write!(f, "folding factor {value} is not currently supported")
            }
            Self::NumPositionEvaluationMismatch(num_positions, num_evaluations) => write!(f,
                "the number of query positions must be the same as the number of polynomial evaluations, but {num_positions} and {num_evaluations} were provided"
            ),
            Self::LayerCommitmentMismatch => {
                write!(f, "FRI queries did not match layer commitment made by the prover")
            }
            Self::InvalidLayerFolding(layer) => {
                write!(f, "degree-respecting projection is not consistent at layer {layer}")
            }
            Self::RemainderCommitmentMismatch => {
                write!(f, "FRI remainder did not match the commitment")
            }
            Self::InvalidRemainderFolding => {
                write!(f, "degree-respecting projection is inconsistent at the last FRI layer")
            }
            Self::RemainderDegreeNotValid => {
                write!(f, "FRI remainder expected degree is greater than number of remainder values")
            }
            Self::RemainderDegreeMismatch(degree) => {
                write!(f, "FRI remainder is not a valid degree {degree} polynomial")
            }
            Self::DegreeTruncation(degree, folding, layer) => {
                write!(f, "degree reduction from {degree} by {folding} at layer {layer} results in degree truncation")
            }
        }
    }
}


// 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 an implementation of the FRI protocol used by the Winterfell STARK prover
//! and verifier.
//!
//! FRI stands for Fast Reed-Solomon Interactive Oracle Proof of Proximity, and is used in the
//! STARK protocol for low-degree testing. Specifically, given a commitment to a set of evaluations
//! of some function over domain *D*, the verifier can be convinced that the function is a
//! polynomial of degree at most *d*, by making a small number of queries to the commitment.
//!
//! # Proof generation
//! FRI proofs are generated by a [FriProver] in two steps:
//!
//! 1. First, the commit phase of the protocol is executed via
//!    [build_layers()](prover::FriProver::build_layers()) function. During this phase, the degree
//!    of the polynomial is repeatedly reduced by applying a degree-respecting projection, until
//!    the size of the domain over which the polynomial is evaluated falls under
//!    `max_remainder_size` parameter. While performing the reduction, the prover writes a set of
//!    layer commitments into the [ProverChannel]. These commitments should be recorded and sent
//!    to the verifier as they will be needed during the proof verification procedure.
//! 2. Then, the query phase of the protocol is executed via
//!    [build_proof()](prover::FriProver::build_proof()) function. The output of this function is
//!    an instance of the [FriProof] struct. When FRI is executed as a part of the STARK protocol,
//!    FRI proof is included into a STARK proof.
//!
//! When the crate is compiled with `concurrent` feature enabled, proof generation will be
//! performed in multiple threads (usually, as many threads as there are logical cores on the
//! machine). The number of threads can be configured via `RAYON_NUM_THREADS` environment variable.
//!
//! # Proof verification
//! FRI proofs are verified by a [FriVerifier] as follows:
//! 1. First, a FRI proof needs to be converted into a [VerifierChannel]. This crate provides a
//!    default implementation of the verifier channel, but when FRI proof verification is executed
//!    as a part of the larger STARK protocol, STARK verifier handles this conversion.
//! 2. Then, a [FriVerifier] should be instantiated (via [new()](FriVerifier::new()) function).
//!    This will execute the commit phase of the FRI protocol from the verifier's perspective -
//!    i.e., the verifier will read FRI layer commitments from the channel, and generates
//!    random values needed for layer folding.
//! 3. Finally, the query phase of the FRI protocol should be executed via
//!    [verify()](FriVerifier::verify()) function. Note that query values at the first FRI layer
//!    are provided to the [verify()](FriVerifier::verify()) function directly. The values at
//!    remaining layers, the verifier reads from the specified verifier channel.
//!
//! # Protocol parameters
//! The current implementation supports executing FRI protocol with dynamically configurable
//! parameters including:
//!
//! * Base STARK field,
//! * Extension field,
//! * Domain blowup factor,
//! * Hash function (used for Merkle tree commitments),
//! * Folding factor (used for degree reduction for each FRI layer),
//! * Maximum size of the last FRI layer.
//!
//! # References
//! * StarkWare's blog post on [Low Degree Testing](https://medium.com/starkware/low-degree-testing-f7614f5172db)
//! * [Fast Reed-Solomon Interactive Oracle Proofs of Proximity](https://eccc.weizmann.ac.il/report/2017/134/)
//! * [DEEP-FRI: Sampling Outside the Box Improves Soundness](https://eprint.iacr.org/2019/336)
//! * Swastik Kooparty's [talk on DEEP-FRI](https://www.youtube.com/watch?v=txo_kPSn59Y&list=PLcIyXLwiPilWvjvNkhMn283LV370Pk5CT&index=6)

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

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

pub mod folding;

mod prover;
pub use prover::{DefaultProverChannel, FriProver, ProverChannel};

mod verifier;
pub use verifier::{DefaultVerifierChannel, FriVerifier, VerifierChannel};

mod options;
pub use options::FriOptions;

mod proof;
pub use proof::FriProof;

mod errors;
pub use errors::VerifierError;

pub mod utils;


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

use math::StarkField;

// FRI OPTIONS
// ================================================================================================

/// FRI protocol config options for proof generation and verification.
#[derive(Clone, PartialEq, Eq)]
pub struct FriOptions {
    folding_factor: usize,
    remainder_max_degree: usize,
    blowup_factor: usize,
}

impl FriOptions {
    /// Returns a new [FriOptions] struct instantiated with the specified parameters.
    ///
    /// # Panics
    /// Panics if:
    /// - `blowup_factor` is not a power of two.
    /// - `folding_factor` is not 2, 4, 8, or 16.
    pub fn new(blowup_factor: usize, folding_factor: usize, remainder_max_degree: usize) -> Self {
        // TODO: change panics to errors
        assert!(
            blowup_factor.is_power_of_two(),
            "blowup factor must be a power of two, but was {blowup_factor}"
        );
        assert!(
            folding_factor == 2
                || folding_factor == 4
                || folding_factor == 8
                || folding_factor == 16,
            "folding factor {folding_factor} is not supported"
        );
        FriOptions {
            folding_factor,
            remainder_max_degree,
            blowup_factor,
        }
    }

    /// Returns the offset by which the evaluation domain is shifted.
    ///
    /// The domain is shifted by multiplying every element in the domain by this offset.
    ///
    /// Currently, the offset is hard-coded to be the primitive element in the field specified by
    /// type parameter `B`.
    pub fn domain_offset<B: StarkField>(&self) -> B {
        B::GENERATOR
    }

    /// Returns the factor by which the degree of a polynomial is reduced with each FRI layer.
    ///
    /// In combination with `remainder_max_degree_plus_1` this property defines how many FRI layers are
    /// needed for an evaluation domain of a given size.
    pub fn folding_factor(&self) -> usize {
        self.folding_factor
    }

    /// Returns maximum allowed remainder polynomial degree.
    ///
    /// In combination with `folding_factor` this property defines how many FRI layers are needed
    /// for an evaluation domain of a given size.
    pub fn remainder_max_degree(&self) -> usize {
        self.remainder_max_degree
    }

    /// Returns a blowup factor of the evaluation domain.
    ///
    /// Specifically, if the polynomial for which the FRI protocol is executed is of degree `d`
    /// where `d` is one less than a power of two, then the evaluation domain size will be
    /// equal to `(d + 1) * blowup_factor`.
    pub fn blowup_factor(&self) -> usize {
        self.blowup_factor
    }

    /// Computes and return the number of FRI layers required for a domain of the specified size.
    ///
    /// The number of layers for a given domain size is defined by the `folding_factor` and
    /// `remainder_max_degree` and `blowup_factor` settings.
    pub fn num_fri_layers(&self, mut domain_size: usize) -> usize {
        let mut result = 0;
        let max_remainder_size = (self.remainder_max_degree + 1) * self.blowup_factor;
        while domain_size > max_remainder_size {
            domain_size /= self.folding_factor;
            result += 1;
        }
        result
    }
}