// 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::errors::AssertionError;
use core::{
    cmp::{Ord, Ordering, PartialOrd},
    fmt::{Display, Formatter},
};
use math::FieldElement;
use utils::collections::Vec;

#[cfg(test)]
mod tests;

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

const MIN_STRIDE_LENGTH: usize = 2;
const NO_STRIDE: usize = 0;

// ASSERTION
// ================================================================================================

/// An assertion made against an execution trace.
///
/// An assertion is always placed against a single column of an execution trace, but can cover
/// multiple steps and multiple values. Specifically, there are three kinds of assertions:
///
/// 1. **Single** assertion - which requires that a value in a single cell of an execution trace
///    is equal to the specified value.
/// 2. **Periodic** assertion - which requires that values in multiple cells of a single column
///   are equal to the specified value. The cells must be evenly spaced at intervals with lengths
///   equal to powers of two. For example, we can specify that values in a column must be equal
///   to 0 at steps 0, 8, 16, 24, 32 etc. Steps can also start at some offset - e.g., 1, 9, 17,
///   25, 33 is also a valid sequence of steps.
/// 3. **Sequence** assertion - which requires that multiple cells in a single column are equal
///   to the values from the provided list. The cells must be evenly spaced at intervals with
///   lengths equal to powers of two. For example, we can specify that values in a column must
///   be equal to a sequence 1, 2, 3, 4 at steps 0, 8, 16, 24. That is, value at step 0 should be
///   equal to 1, value at step 8 should be equal to 2 etc.
///
/// Note that single and periodic assertions are succinct. That is, a verifier can evaluate them
/// very efficiently. However, sequence assertions have liner complexity in the number of
/// asserted values. Though, unless many thousands of values are asserted, practical impact of
/// this linear complexity should be negligible.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Assertion<E: FieldElement> {
    pub(super) column: usize,
    pub(super) first_step: usize,
    pub(super) stride: usize,
    pub(super) values: Vec<E>,
}

impl<E: FieldElement> Assertion<E> {
    // CONSTRUCTORS
    // --------------------------------------------------------------------------------------------
    /// Returns an assertion against a single cell of an execution trace.
    ///
    /// The returned assertion requires that the value in the specified `column` at the specified
    /// `step` is equal to the provided `value`.
    pub fn single(column: usize, step: usize, value: E) -> Self {
        Assertion {
            column,
            first_step: step,
            stride: NO_STRIDE,
            values: vec![value],
        }
    }

    /// Returns an single-value assertion against multiple cells of a single column.
    ///
    /// The returned assertion requires that values in the specified `column` must be equal to
    /// the specified `value` at steps which start at `first_step` and repeat in equal intervals
    /// specified by `stride`.
    ///
    /// # Panics
    /// Panics if:
    /// * `stride` is not a power of two, or is smaller than 2.
    /// * `first_step` is greater than `stride`.
    pub fn periodic(column: usize, first_step: usize, stride: usize, value: E) -> Self {
        validate_stride(stride, first_step, column);
        Assertion {
            column,
            first_step,
            stride,
            values: vec![value],
        }
    }

    /// Returns a multi-value assertion against multiple cells of a single column.
    ///
    /// The returned assertion requires that values in the specified `column` must be equal to
    /// the provided `values` at steps which start at `first_step` and repeat in equal intervals
    /// specified by `stride` until all values have been consumed.
    ///
    /// # Panics
    /// Panics if:
    /// * `stride` is not a power of two, or is smaller than 2.
    /// * `first_step` is greater than `stride`.
    /// * `values` is empty or number of values in not a power of two.
    pub fn sequence(column: usize, first_step: usize, stride: usize, values: Vec<E>) -> Self {
        validate_stride(stride, first_step, column);
        assert!(
            !values.is_empty(),
            "invalid assertion for column {column}: number of asserted values must be greater than zero"
        );
        assert!(
            values.len().is_power_of_two(),
            "invalid assertion for column {}: number of asserted values must be a power of two, but was {}",
            column,
            values.len()
        );
        Assertion {
            column,
            first_step,
            stride: if values.len() == 1 { NO_STRIDE } else { stride },
            values,
        }
    }

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

    /// Returns index of the column against which this assertion is placed.
    pub fn column(&self) -> usize {
        self.column
    }

    /// Returns the first step of the execution trace against which this assertion is placed.
    ///
    /// For single value assertions this is equivalent to the assertion step.
    pub fn first_step(&self) -> usize {
        self.first_step
    }

    /// Returns the interval at which the assertion repeats in the execution trace.
    ///
    /// For single value assertions, this will be 0.
    pub fn stride(&self) -> usize {
        self.stride
    }

    /// Returns asserted values.
    ///
    /// For single value and periodic assertions this will be a slice containing one value.
    pub fn values(&self) -> &[E] {
        &self.values
    }

    /// Returns true if this is a single-value assertion (one value, one step).
    pub fn is_single(&self) -> bool {
        self.stride == NO_STRIDE
    }

    /// Returns true if this is a periodic assertion (one value, many steps).
    pub fn is_periodic(&self) -> bool {
        self.stride != NO_STRIDE && self.values.len() == 1
    }

    /// Returns true if this is a sequence assertion (many values, many steps).
    pub fn is_sequence(&self) -> bool {
        self.values.len() > 1
    }

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

    /// Checks if this assertion overlaps with the provided assertion.
    ///
    /// Overlap is defined as asserting a value for the same step in the same column.
    pub fn overlaps_with(&self, other: &Assertion<E>) -> bool {
        if self.column != other.column {
            return false;
        }
        if self.first_step == other.first_step {
            return true;
        }
        if self.stride == other.stride {
            return false;
        }

        // at this point we know that assertions are for the same column but they start
        // on different steps and also have different strides

        if self.first_step < other.first_step {
            if self.is_single() {
                return false;
            }
            if other.is_single() || self.stride < other.stride {
                (other.first_step - self.first_step) % self.stride == 0
            } else {
                false
            }
        } else {
            if other.is_single() {
                return false;
            }
            if self.is_single() || other.stride < self.stride {
                (self.first_step - other.first_step) % other.stride == 0
            } else {
                false
            }
        }
    }

    /// Panics if the assertion cannot be placed against an execution trace of the specified width.
    pub fn validate_trace_width(&self, trace_width: usize) -> Result<(), AssertionError> {
        if self.column >= trace_width {
            return Err(AssertionError::TraceWidthTooShort(self.column, trace_width));
        }
        Ok(())
    }

    /// Checks if the assertion is valid against an execution trace of the specified length.
    ///
    /// # Errors
    /// Returns an error if:
    /// * `trace_length` is not a power of two.
    /// * For single assertion, `first_step` >= `trace_length`.
    /// * For periodic assertion, `stride` > `trace_length`.
    /// * For sequence assertion, `num_values` * `stride` != `trace_length`;
    pub fn validate_trace_length(&self, trace_length: usize) -> Result<(), AssertionError> {
        if !trace_length.is_power_of_two() {
            return Err(AssertionError::TraceLengthNotPowerOfTwo(trace_length));
        }
        if self.is_single() {
            if self.first_step >= trace_length {
                return Err(AssertionError::TraceLengthTooShort(
                    (self.first_step + 1).next_power_of_two(),
                    trace_length,
                ));
            }
        } else if self.is_periodic() {
            if self.stride > trace_length {
                return Err(AssertionError::TraceLengthTooShort(
                    self.stride,
                    trace_length,
                ));
            }
        } else {
            let expected_length = self.values.len() * self.stride;
            if expected_length != trace_length {
                return Err(AssertionError::TraceLengthNotExact(
                    expected_length,
                    trace_length,
                ));
            }
        }
        Ok(())
    }

    /// Executes the provided closure for all possible instantiations of this assertions against
    /// a execution trace of the specified length.
    ///
    /// # Panics
    /// Panics if the specified trace length is not valid for this assertion.
    pub fn apply<F>(&self, trace_length: usize, mut f: F)
    where
        F: FnMut(usize, E),
    {
        self.validate_trace_length(trace_length)
            .unwrap_or_else(|err| {
                panic!("invalid trace length: {err}");
            });
        if self.is_single() {
            f(self.first_step, self.values[0]);
        } else if self.is_periodic() {
            for i in 0..(trace_length / self.stride) {
                f(self.first_step + self.stride * i, self.values[0]);
            }
        } else {
            for (i, &value) in self.values.iter().enumerate() {
                f(self.first_step + self.stride * i, value);
            }
        }
    }

    /// Returns the number of steps against which this assertion will be applied given an
    /// execution trace of the specified length.
    ///
    /// * For single-value assertions, this will always be one.
    /// * For periodic assertions this will be equal to `trace_length` / `stride`.
    /// * For sequence assertions this will be equal to the number of asserted values.
    ///
    /// # Panics
    /// Panics if the specified trace length is not valid for this assertion.
    pub fn get_num_steps(&self, trace_length: usize) -> usize {
        self.validate_trace_length(trace_length)
            .unwrap_or_else(|err| {
                panic!("invalid trace length: {err}");
            });
        if self.is_single() {
            1
        } else if self.is_periodic() {
            trace_length / self.stride
        } else {
            self.values.len()
        }
    }
}

// OTHER TRAIT IMPLEMENTATIONS
// =================================================================================================

/// We define ordering of assertions to be first by stride, then by first_step, and finally by
/// column in ascending order.
impl<E: FieldElement> Ord for Assertion<E> {
    fn cmp(&self, other: &Self) -> Ordering {
        if self.stride == other.stride {
            if self.first_step == other.first_step {
                self.column.partial_cmp(&other.column).unwrap()
            } else {
                self.first_step.partial_cmp(&other.first_step).unwrap()
            }
        } else {
            self.stride.partial_cmp(&other.stride).unwrap()
        }
    }
}

impl<E: FieldElement> PartialOrd for Assertion<E> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<E: FieldElement> Display for Assertion<E> {
    fn fmt(&self, f: &mut Formatter) -> core::fmt::Result {
        write!(f, "(column={}, ", self.column)?;
        match self.stride {
            0 => write!(f, "step={}, ", self.first_step)?,
            _ => {
                let second_step = self.first_step + self.stride;
                write!(f, "steps=[{}, {}, ...], ", self.first_step, second_step)?;
            }
        }
        match self.values.len() {
            1 => write!(f, "value={})", self.values[0]),
            2 => write!(f, "values=[{}, {}])", self.values[0], self.values[1]),
            _ => write!(f, "values=[{}, {}, ...])", self.values[0], self.values[1]),
        }
    }
}

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

fn validate_stride(stride: usize, first_step: usize, column: usize) {
    assert!(
        stride.is_power_of_two(),
        "invalid assertion for column {column}: stride must be a power of two, but was {stride}"
    );
    assert!(
        stride >= MIN_STRIDE_LENGTH,
        "invalid assertion for column {column}: stride must be at least {MIN_STRIDE_LENGTH}, but was {stride}"
    );
    assert!(
        first_step < stride,
        "invalid assertion for column {column}: first step must be smaller than stride ({stride} steps), but was {first_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::{Assertion, AssertionError};
use math::{fields::f128::BaseElement, FieldElement};
use rand_utils::{rand_value, rand_vector};
use utils::collections::Vec;

// SINGLE ASSERTIONS
// ================================================================================================
#[test]
fn single_assertion() {
    let value = rand_value::<BaseElement>();
    let a = Assertion::single(2, 8, value);
    assert_eq!(2, a.column);
    assert_eq!(8, a.first_step);
    assert_eq!(vec![value], a.values);
    assert_eq!(0, a.stride);
    assert_eq!(1, a.get_num_steps(16));
    assert_eq!(1, a.get_num_steps(32));

    a.apply(16, |step, val| {
        assert_eq!(8, step);
        assert_eq!(value, val);
    });

    assert_eq!(Ok(()), a.validate_trace_width(3));
    assert_eq!(
        Err(AssertionError::TraceWidthTooShort(2, 1)),
        a.validate_trace_width(1)
    );

    assert_eq!(Ok(()), a.validate_trace_length(16));
    assert_eq!(
        Err(AssertionError::TraceLengthTooShort(16, 8)),
        a.validate_trace_length(8)
    );
}

// PERIODIC ASSERTIONS
// ================================================================================================

#[test]
fn periodic_assertion() {
    let value = rand_value::<BaseElement>();
    let a = Assertion::periodic(4, 1, 16, value);
    assert_eq!(4, a.column);
    assert_eq!(1, a.first_step);
    assert_eq!(vec![value], a.values);
    assert_eq!(16, a.stride);
    assert_eq!(1, a.get_num_steps(16));
    assert_eq!(2, a.get_num_steps(32));

    a.apply(16, |step, val| {
        assert_eq!(1, step);
        assert_eq!(value, val);
    });
    a.apply(32, |step, val| {
        if step == 1 || step == 17 {
            assert_eq!(value, val);
            return;
        }
        unreachable!();
    });

    assert_eq!(Ok(()), a.validate_trace_width(5));
    assert_eq!(
        Err(AssertionError::TraceWidthTooShort(4, 2)),
        a.validate_trace_width(2)
    );
    assert_eq!(Ok(()), a.validate_trace_length(16));
    assert_eq!(
        Err(AssertionError::TraceLengthTooShort(16, 8)),
        a.validate_trace_length(8)
    );
}

#[test]
#[should_panic(
    expected = "invalid assertion for column 0: stride must be a power of two, but was 3"
)]
fn periodic_assertion_stride_not_power_of_two() {
    let _ = Assertion::periodic(0, 1, 3, BaseElement::ONE);
}

#[test]
#[should_panic(expected = "invalid assertion for column 0: stride must be at least 2, but was 1")]
fn periodic_assertion_stride_too_small() {
    let _ = Assertion::periodic(0, 1, 1, BaseElement::ONE);
}

#[test]
#[should_panic(
    expected = "invalid assertion for column 0: first step must be smaller than stride (4 steps), but was 5"
)]
fn periodic_assertion_first_step_greater_than_stride() {
    let _ = Assertion::periodic(0, 5, 4, BaseElement::ONE);
}

#[test]
#[should_panic(
    expected = "invalid trace length: expected trace length to be at least 8, but was 4"
)]
fn periodic_assertion_get_num_steps_error() {
    let a = Assertion::periodic(0, 1, 8, BaseElement::ONE);
    let _ = a.get_num_steps(4);
}

// SEQUENCE ASSERTIONS
// ================================================================================================

#[test]
fn sequence_assertion() {
    let values = rand_vector::<BaseElement>(2);
    let a = Assertion::sequence(3, 2, 4, values.clone());
    assert_eq!(3, a.column);
    assert_eq!(2, a.first_step);
    assert_eq!(values, a.values);
    assert_eq!(4, a.stride);
    assert_eq!(2, a.get_num_steps(8));

    a.apply(8, |step, val| {
        if step == 2 {
            assert_eq!(values[0], val);
            return;
        } else if step == 6 {
            assert_eq!(values[1], val);
            return;
        }
        unreachable!();
    });

    assert_eq!(Ok(()), a.validate_trace_length(8));
    assert_eq!(
        Err(AssertionError::TraceLengthNotExact(8, 4)),
        a.validate_trace_length(4)
    );
    assert_eq!(
        Err(AssertionError::TraceLengthNotExact(8, 16)),
        a.validate_trace_length(16)
    );

    assert_eq!(Ok(()), a.validate_trace_width(4));
    assert_eq!(
        Err(AssertionError::TraceWidthTooShort(3, 2)),
        a.validate_trace_width(2)
    );
}

#[test]
#[should_panic(
    expected = "invalid assertion for column 3: stride must be a power of two, but was 5"
)]
fn sequence_assertion_stride_not_power_of_two() {
    let _ = Assertion::sequence(3, 2, 5, vec![BaseElement::ONE, BaseElement::ZERO]);
}

#[test]
#[should_panic(expected = "invalid assertion for column 3: stride must be at least 2, but was 1")]
fn sequence_assertion_stride_too_small() {
    let _ = Assertion::sequence(3, 2, 1, vec![BaseElement::ONE, BaseElement::ZERO]);
}

#[test]
#[should_panic(
    expected = "invalid assertion for column 3: first step must be smaller than stride (4 steps), but was 5"
)]
fn sequence_assertion_first_step_greater_than_stride() {
    let _ = Assertion::sequence(3, 5, 4, vec![BaseElement::ONE, BaseElement::ZERO]);
}

#[test]
#[should_panic(expected = "invalid trace length: expected trace length to be exactly 8, but was 4")]
fn sequence_assertion_inconsistent_trace() {
    let a = Assertion::sequence(3, 2, 4, vec![BaseElement::ONE, BaseElement::ZERO]);
    let _ = a.get_num_steps(4);
}

#[test]
#[should_panic(
    expected = "invalid assertion for column 3: number of asserted values must be greater than zero"
)]
fn sequence_assertion_empty_values() {
    let _ = Assertion::sequence(3, 2, 4, Vec::<BaseElement>::new());
}

#[test]
#[should_panic(
    expected = "invalid assertion for column 3: number of asserted values must be a power of two, but was 3"
)]
fn sequence_assertion_num_values_not_power_of_two() {
    let _ = Assertion::sequence(
        3,
        2,
        4,
        vec![BaseElement::ONE, BaseElement::ZERO, BaseElement::ONE],
    );
}

// OVERLAPPING ASSERTIONS
// ================================================================================================

#[test]
fn assertion_overlap() {
    // ----- single-single overlap ----------------------------------------------------------------

    let a = Assertion::single(3, 2, BaseElement::ONE);
    let b = Assertion::single(3, 2, BaseElement::ONE);
    assert!(a.overlaps_with(&b));

    // different columns: no overlap
    let b = Assertion::single(1, 2, BaseElement::ONE);
    assert!(!a.overlaps_with(&b));

    // different steps: no overlap
    let b = Assertion::single(3, 1, BaseElement::ONE);
    assert!(!a.overlaps_with(&b));

    // ----- single-periodic overlap --------------------------------------------------------------

    let a = Assertion::periodic(3, 2, 4, BaseElement::ONE);
    let b = Assertion::single(3, 2, BaseElement::ONE);
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    let b = Assertion::single(3, 6, BaseElement::ONE);
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    let b = Assertion::single(3, 10, BaseElement::ONE);
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    // different columns: no overlap
    let b = Assertion::single(1, 2, BaseElement::ONE);
    assert!(!a.overlaps_with(&b));
    assert!(!b.overlaps_with(&a));

    // different steps: no overlap
    let b = Assertion::single(3, 3, BaseElement::ONE);
    assert!(!a.overlaps_with(&b));
    assert!(!b.overlaps_with(&a));

    // ----- single-sequence overlap --------------------------------------------------------------

    let values = vec![BaseElement::ONE, BaseElement::ZERO];
    let a = Assertion::sequence(3, 2, 8, values);
    let b = Assertion::single(3, 2, BaseElement::ONE);
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    let b = Assertion::single(3, 10, BaseElement::ONE);
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    let b = Assertion::single(3, 18, BaseElement::ONE);
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    // different columns: no overlap
    let b = Assertion::single(1, 2, BaseElement::ONE);
    assert!(!a.overlaps_with(&b));
    assert!(!b.overlaps_with(&a));

    // different steps: no overlap
    let b = Assertion::single(3, 3, BaseElement::ONE);
    assert!(!a.overlaps_with(&b));
    assert!(!b.overlaps_with(&a));

    // ----- periodic-periodic overlap ------------------------------------------------------------

    let a = Assertion::periodic(3, 4, 8, BaseElement::ONE);
    let b = Assertion::periodic(3, 4, 8, BaseElement::ONE);
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    let b = Assertion::periodic(3, 4, 16, BaseElement::ONE);
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    let b = Assertion::periodic(3, 0, 4, BaseElement::ONE);
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    // different columns: no overlap
    let b = Assertion::periodic(1, 4, 8, BaseElement::ONE);
    assert!(!a.overlaps_with(&b));
    assert!(!b.overlaps_with(&a));

    // different first step: no overlap
    let b = Assertion::periodic(0, 0, 8, BaseElement::ONE);
    assert!(!a.overlaps_with(&b));
    assert!(!b.overlaps_with(&a));

    // different first step and bigger stride: no overlap
    let b = Assertion::periodic(0, 0, 16, BaseElement::ONE);
    assert!(!a.overlaps_with(&b));
    assert!(!b.overlaps_with(&a));

    // ----- sequence-sequence overlap ------------------------------------------------------------

    let values = vec![BaseElement::ONE, BaseElement::ZERO];

    let a = Assertion::sequence(3, 4, 8, values.clone());
    let b = Assertion::sequence(3, 4, 8, values.clone());
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    let b = Assertion::sequence(3, 4, 16, values.clone());
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    let b = Assertion::sequence(3, 0, 4, values.clone());
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    // different columns: no overlap
    let b = Assertion::sequence(1, 4, 8, values.clone());
    assert!(!a.overlaps_with(&b));
    assert!(!b.overlaps_with(&a));

    // different first step: no overlap
    let b = Assertion::sequence(0, 0, 8, values.clone());
    assert!(!a.overlaps_with(&b));
    assert!(!b.overlaps_with(&a));

    // different first step and bigger stride: no overlap
    let b = Assertion::sequence(0, 0, 16, values);
    assert!(!a.overlaps_with(&b));
    assert!(!b.overlaps_with(&a));

    // ----- sequence-periodic overlap ------------------------------------------------------------

    let values = vec![BaseElement::ONE, BaseElement::ZERO];

    let a = Assertion::sequence(3, 4, 8, values);
    let b = Assertion::periodic(3, 4, 8, BaseElement::ONE);
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    let b = Assertion::periodic(3, 4, 16, BaseElement::ONE);
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    let b = Assertion::periodic(3, 0, 4, BaseElement::ONE);
    assert!(a.overlaps_with(&b));
    assert!(b.overlaps_with(&a));

    // different columns: no overlap
    let b = Assertion::periodic(1, 4, 8, BaseElement::ONE);
    assert!(!a.overlaps_with(&b));
    assert!(!b.overlaps_with(&a));

    // different first step: no overlap
    let b = Assertion::periodic(0, 0, 8, BaseElement::ONE);
    assert!(!a.overlaps_with(&b));
    assert!(!b.overlaps_with(&a));

    // different first step and bigger stride: no overlap
    let b = Assertion::periodic(0, 0, 16, BaseElement::ONE);
    assert!(!a.overlaps_with(&b));
    assert!(!b.overlaps_with(&a));
}


// 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::{Assertion, BTreeMap, ExtensionOf, FieldElement, Vec};
use math::{fft, polynom};

// BOUNDARY CONSTRAINT
// ================================================================================================
/// The numerator portion of a boundary constraint.
///
/// A boundary constraint is described by a rational function $\frac{f(x) - b(x)}{z(x)}$, where:
///
/// * $f(x)$ is a trace polynomial for the column against which the constraint is placed.
/// * $b(b)$ is the value polynomial for this constraint.
/// * $z(x)$ is the constraint divisor polynomial.
///
/// In addition to the value polynomial, a [BoundaryConstraint] also contains info needed to
/// evaluate the constraint and to compose constraint evaluations with other constraints (i.e.,
/// constraint composition coefficient).
///
/// When the protocol is run in a large field, types `F` and `E` are the same. However, when
/// working with small fields, `F` and `E` can be set as follows:
/// * `F` could be the base field of the protocol, in which case `E` is the extension field used.
/// * `F` could be the extension field, in which case `F` and `E` are the same type.
///
/// Boundary constraints cannot be instantiated directly, they are created internally from
/// [Assertions](Assertion).
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct BoundaryConstraint<F, E>
where
    F: FieldElement,
    E: FieldElement<BaseField = F::BaseField> + ExtensionOf<F>,
{
    column: usize,
    poly: Vec<F>,
    poly_offset: (usize, F::BaseField),
    cc: E,
}

impl<F, E> BoundaryConstraint<F, E>
where
    F: FieldElement,
    E: FieldElement<BaseField = F::BaseField> + ExtensionOf<F>,
{
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Creates a new boundary constraint from the specified assertion.
    pub(super) fn new(
        assertion: Assertion<F>,
        inv_g: F::BaseField,
        twiddle_map: &mut BTreeMap<usize, Vec<F::BaseField>>,
        composition_coefficient: E,
    ) -> Self {
        // build a polynomial which evaluates to constraint values at asserted steps; for
        // single-value assertions we use the value as constant coefficient of degree 0
        // polynomial; but for multi-value assertions, we need to interpolate the values
        // into a polynomial using inverse FFT
        let mut poly_offset = (0, F::BaseField::ONE);
        let mut poly = assertion.values;
        if poly.len() > 1 {
            // get the twiddles from the map; if twiddles for this domain haven't been built
            // yet, build them and add them to the map
            let inv_twiddles = twiddle_map
                .entry(poly.len())
                .or_insert_with(|| fft::get_inv_twiddles(poly.len()));
            // interpolate the values into a polynomial
            fft::interpolate_poly(&mut poly, inv_twiddles);
            if assertion.first_step != 0 {
                // if the assertions don't fall on the steps which are powers of two, we can't
                // use FFT to interpolate the values into a polynomial. This would make such
                // assertions quite impractical. To get around this, we still use FFT to build
                // the polynomial, but then we evaluate it as f(x * offset) instead of f(x)
                let x_offset = inv_g.exp((assertion.first_step as u64).into());
                poly_offset = (assertion.first_step, x_offset);
            }
        }

        BoundaryConstraint {
            column: assertion.column,
            poly,
            poly_offset,
            cc: composition_coefficient,
        }
    }

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

    /// Returns index of the column against which this constraint applies.
    pub fn column(&self) -> usize {
        self.column
    }

    /// Returns a value polynomial for this constraint.
    pub fn poly(&self) -> &[F] {
        &self.poly
    }

    /// Returns offset by which we need to shift the domain before evaluating this constraint.
    ///
    /// The offset is returned as a tuple describing both, the number of steps by which the
    /// domain needs to be shifted, and field element by which a domain element needs to be
    /// multiplied to achieve the desired shift.
    pub fn poly_offset(&self) -> (usize, F::BaseField) {
        self.poly_offset
    }

    /// Returns composition coefficient for this constraint.
    pub fn cc(&self) -> &E {
        &self.cc
    }

    // CONSTRAINT EVALUATOR
    // --------------------------------------------------------------------------------------------
    /// Evaluates this constraint at the specified point `x`.
    ///
    /// The constraint is evaluated by computing $f(x) - b(x)$, where:
    /// * $f$ is a trace polynomial for the column against which the constraint is placed.
    /// * $f(x)$ = `trace_value`
    /// * $b$ is the value polynomial for this constraint.
    ///
    /// For boundary constraints derived from single and periodic assertions, $b(x)$ is a constant.
    pub fn evaluate_at(&self, x: E, trace_value: E) -> E {
        let assertion_value = if self.poly.len() == 1 {
            // if the value polynomial consists of just a constant, use that constant
            E::from(self.poly[0])
        } else {
            // otherwise, we need to evaluate the polynomial at `x`; for assertions which don't
            // fall on steps that are powers of two, we need to evaluate the value polynomial
            // at x * offset (instead of just x).
            //
            // note that while the coefficients of the value polynomial are in the base field,
            // if we are working in an extension field, the result of the evaluation will be a
            // value in the extension field.
            let x = x * E::from(self.poly_offset.1);
            polynom::eval(&self.poly, x)
        };
        // subtract assertion value from trace value
        trace_value - assertion_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.

use super::{
    Assertion, BTreeMap, BoundaryConstraint, ConstraintDivisor, ExtensionOf, FieldElement, Vec,
};

// BOUNDARY CONSTRAINT GROUP
// ================================================================================================
/// A group of boundary constraints all having the same divisor.
///
/// A boundary constraint is described by a rational function $\frac{f(x) - b(x)}{z(x)}$, where:
///
/// * $f(x)$ is a trace polynomial for the column against which the constraint is placed.
/// * $b(x)$ is the value polynomial for the constraint.
/// * $z(x)$ is the constraint divisor polynomial.
///
/// A boundary constraint group groups together all boundary constraints where polynomial $z$ is
/// the same. The constraints stored in the group describe polynomials $b$. At the time of
/// constraint evaluation, a prover or a verifier provides evaluations of the relevant polynomial
/// $f$ so that the value of the constraint can be computed.
///
/// When the protocol is run in a large field, types `F` and `E` are the same. However, when
/// working with small fields, `F` and `E` can be set as follows:
/// * `F` could be the base field of the protocol, in which case `E` is the extension field used.
/// * `F` could be the extension field, in which case `F` and `E` are the same type.
///
/// The above arrangement allows us to describe boundary constraints for main and auxiliary
/// segments of the execution trace. Specifically:
/// * For the constraints against columns of the main execution trace, `F` is set to the base field
///   of the protocol, and `E` is set to the extension field.
/// * For the constraints against columns of auxiliary trace segments, both `F` and `E` are set to
///   the extension field.
#[derive(Debug, Clone)]
pub struct BoundaryConstraintGroup<F, E>
where
    F: FieldElement,
    E: FieldElement<BaseField = F::BaseField> + ExtensionOf<F>,
{
    constraints: Vec<BoundaryConstraint<F, E>>,
    divisor: ConstraintDivisor<F::BaseField>,
}

impl<F, E> BoundaryConstraintGroup<F, E>
where
    F: FieldElement,
    E: FieldElement<BaseField = F::BaseField> + ExtensionOf<F>,
{
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new boundary constraint group to hold constraints with the specified divisor.
    pub(super) fn new(divisor: ConstraintDivisor<F::BaseField>) -> Self {
        BoundaryConstraintGroup {
            constraints: Vec::new(),
            divisor,
        }
    }

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

    /// Returns a list of boundary constraints in this group.
    pub fn constraints(&self) -> &[BoundaryConstraint<F, E>] {
        &self.constraints
    }

    /// Returns a divisor applicable to all boundary constraints in this group.
    pub fn divisor(&self) -> &ConstraintDivisor<F::BaseField> {
        &self.divisor
    }

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

    /// Creates a new boundary constraint from the specified assertion and adds it to the group.
    pub(super) fn add(
        &mut self,
        assertion: Assertion<F>,
        inv_g: F::BaseField,
        twiddle_map: &mut BTreeMap<usize, Vec<F::BaseField>>,
        composition_coefficients: E,
    ) {
        self.constraints.push(BoundaryConstraint::new(
            assertion,
            inv_g,
            twiddle_map,
            composition_coefficients,
        ));
    }

    /// Evaluates all constraints in this group at the specified point `x`.
    ///
    /// Constraint evaluations are merges into a single value by computing their random linear
    /// combination and dividing the result by the divisor of this constraint group as follows:
    /// $$
    /// \frac{\sum_{i=0}^{k-1}{\alpha_i \cdot C_i(x)}}{z(x)}
    /// $$
    /// where:
    /// * $C_i(x)$ is the evaluation of the $i$th constraint at `x` computed as $f(x) - b(x)$.
    /// * $\alpha_i$ are random field elements. In the interactive version of the
    ///   protocol, these are provided by the verifier.
    pub fn evaluate_at(&self, state: &[E], x: E) -> E {
        let mut numerator = E::ZERO;
        for constraint in self.constraints().iter() {
            let trace_value = state[constraint.column()];
            let evaluation = constraint.evaluate_at(x, trace_value);
            numerator += evaluation * *constraint.cc();
        }

        let denominator = self.divisor.evaluate_at(x);

        numerator / denominator
    }
}


// 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::{AirContext, Assertion, ConstraintDivisor};
use math::{ExtensionOf, FieldElement};
use utils::collections::{BTreeMap, BTreeSet, Vec};

mod constraint;
pub use constraint::BoundaryConstraint;

mod constraint_group;
pub use constraint_group::BoundaryConstraintGroup;

#[cfg(test)]
mod tests;

// BOUNDARY CONSTRAINT INFO
// ================================================================================================

/// Boundary constraints for a computation.
///
/// Boundary constraints are arranged into two categories: constraints against columns of the main
/// trace segment, and constraints against columns of auxiliary trace segments. Within each
/// category, the constraints are grouped by their divisor (see [BoundaryConstraintGroup] for
/// more info on boundary constraint structure).
///
/// When the protocol is run in a large field, types `B` and `E` are the same. However, when
/// working with small fields, these types are used as follows:
/// * Constraints against columns of the main segment of the execution trace are defined over the
///   field specified by `B`.
/// * Constraints against columns of auxiliary segments of the execution trace (if any) are defined
///   over the field specified by `E`.
/// * Constraint composition coefficients are defined over the field specified by `E`.
/// * Constraint divisors are defined over the field specified by `B`.
pub struct BoundaryConstraints<E: FieldElement> {
    main_constraints: Vec<BoundaryConstraintGroup<E::BaseField, E>>,
    aux_constraints: Vec<BoundaryConstraintGroup<E, E>>,
}

impl<E: FieldElement> BoundaryConstraints<E> {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new instance of [BoundaryConstraints] for a computation described by the provided
    /// assertions and AIR context.
    ///
    /// # Panics
    /// Panics if:
    /// * The number of provided assertions does not match the number of assertions described by
    ///   the context.
    /// * The number of assertions does not match the number of the provided composition
    ///   coefficients.
    /// * The specified assertions are not valid in the context of the computation (e.g., assertion
    ///   column index is out of bounds).
    pub fn new(
        context: &AirContext<E::BaseField>,
        main_assertions: Vec<Assertion<E::BaseField>>,
        aux_assertions: Vec<Assertion<E>>,
        composition_coefficients: &[E],
    ) -> Self {
        // make sure the provided assertions are consistent with the specified context
        assert_eq!(
            main_assertions.len(),
            context.num_main_assertions,
            "expected {} assertions against main trace segment, but received {}",
            context.num_main_assertions,
            main_assertions.len(),
        );

        assert_eq!(
            aux_assertions.len(),
            context.num_aux_assertions,
            "expected {} assertions against auxiliary trace segments, but received {}",
            context.num_aux_assertions,
            aux_assertions.len(),
        );

        assert_eq!(
            context.num_assertions(),
            composition_coefficients.len(),
            "number of assertions must match the number of composition coefficient tuples"
        );

        let trace_length = context.trace_info.length();
        let main_trace_width = context.trace_info.layout().main_trace_width();
        let aux_trace_width = context.trace_info.layout().aux_trace_width();

        // make sure the assertions are valid in the context of their respective trace segments;
        // also, sort the assertions in the deterministic order so that changing the order of
        // assertions does not change random coefficients that get assigned to them.
        let main_assertions = prepare_assertions(main_assertions, main_trace_width, trace_length);
        let aux_assertions = prepare_assertions(aux_assertions, aux_trace_width, trace_length);

        // compute inverse of the trace domain generator; this will be used for offset
        // computations when creating sequence constraints
        let inv_g = context.trace_domain_generator.inv();

        // cache inverse twiddles for multi-value assertions in this map so that we don't have
        // to re-build them for assertions with identical strides
        let mut twiddle_map = BTreeMap::new();

        // split composition coefficients into main and auxiliary parts
        let (main_composition_coefficients, aux_composition_coefficients) =
            composition_coefficients.split_at(main_assertions.len());

        // build constraints for the assertions against the main trace segment
        let main_constraints = group_constraints(
            main_assertions,
            context,
            main_composition_coefficients,
            inv_g,
            &mut twiddle_map,
        );

        // build constraints for the assertions against auxiliary trace segments
        let aux_constraints = group_constraints(
            aux_assertions,
            context,
            aux_composition_coefficients,
            inv_g,
            &mut twiddle_map,
        );

        Self {
            main_constraints,
            aux_constraints,
        }
    }

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

    /// Returns a reference to the boundary constraints against the main segment of an execution
    /// trace. The constraints are grouped by their divisors.
    pub fn main_constraints(&self) -> &[BoundaryConstraintGroup<E::BaseField, E>] {
        &self.main_constraints
    }

    /// Returns a reference to the boundary constraints against auxiliary segments of an execution
    /// trace. The constraints are grouped by their divisors.
    pub fn aux_constraints(&self) -> &[BoundaryConstraintGroup<E, E>] {
        &self.aux_constraints
    }
}

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

/// Translates the provided assertions into boundary constraints, groups the constraints by their
/// divisor, and sorts the resulting groups by the degree adjustment factor.
fn group_constraints<F, E>(
    assertions: Vec<Assertion<F>>,
    context: &AirContext<F::BaseField>,
    composition_coefficients: &[E],
    inv_g: F::BaseField,
    twiddle_map: &mut BTreeMap<usize, Vec<F::BaseField>>,
) -> Vec<BoundaryConstraintGroup<F, E>>
where
    F: FieldElement,
    E: FieldElement<BaseField = F::BaseField> + ExtensionOf<F>,
{
    // iterate over all assertions, which are sorted first by stride and then by first_step
    // in ascending order
    let mut groups = BTreeMap::new();
    for (assertion, &cc) in assertions.into_iter().zip(composition_coefficients) {
        let key = (assertion.stride(), assertion.first_step());
        let group = groups.entry(key).or_insert_with(|| {
            BoundaryConstraintGroup::new(ConstraintDivisor::from_assertion(
                &assertion,
                context.trace_len(),
            ))
        });

        // add a new assertion constraint to the current group (last group in the list)
        group.add(assertion, inv_g, twiddle_map, cc);
    }

    //return a vector of groups
    groups.into_iter().map(|e| e.1).collect::<Vec<_>>()
}

/// Makes sure the assertions are valid in the context of this computation and don't overlap with
/// each other - i.e. no two assertions are placed against the same column and step combination.
///
/// This also sorts the assertions in their 'natural order'. The natural order is defined as
/// sorting first by stride, then by first step, and finally by column, all in ascending order.
fn prepare_assertions<E: FieldElement>(
    assertions: Vec<Assertion<E>>,
    trace_width: usize,
    trace_length: usize,
) -> Vec<Assertion<E>> {
    // we use a sorted set to help us sort the assertions by their 'natural' order
    let mut result = BTreeSet::<Assertion<E>>::new();

    for assertion in assertions.into_iter() {
        assertion
            .validate_trace_width(trace_width)
            .unwrap_or_else(|err| {
                panic!("assertion {assertion} is invalid: {err}");
            });
        assertion
            .validate_trace_length(trace_length)
            .unwrap_or_else(|err| {
                panic!("assertion {assertion} is invalid: {err}");
            });
        for a in result.iter().filter(|a| a.column == assertion.column) {
            assert!(
                !a.overlaps_with(&assertion),
                "assertion {assertion} overlaps with assertion {a}"
            );
        }

        result.insert(assertion);
    }

    result.into_iter().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 super::{
    super::tests::{build_prng, build_sequence_poly},
    Assertion, BoundaryConstraint,
};
use crypto::{hashers::Blake3_256, DefaultRandomCoin, RandomCoin};
use math::{fields::f64::BaseElement, polynom, FieldElement, StarkField};
use rand_utils::{rand_value, rand_vector, shuffle};
use utils::collections::{BTreeMap, Vec};

// BOUNDARY CONSTRAINT TESTS
// ================================================================================================

#[test]
fn boundary_constraint_from_single_assertion() {
    let mut test_prng = build_prng();
    let (inv_g, mut twiddle_map, mut prng) = build_constraint_params(16);

    // constraint should be built correctly for column 0, step 0
    let value = rand_value::<BaseElement>();
    let assertion = Assertion::single(0, 0, value);
    let constraint = BoundaryConstraint::<BaseElement, BaseElement>::new(
        assertion,
        inv_g,
        &mut twiddle_map,
        prng.draw().unwrap(),
    );
    assert_eq!(0, constraint.column());
    assert_eq!(vec![value], constraint.poly());
    assert_eq!((0, BaseElement::ONE), constraint.poly_offset());
    assert_eq!(&test_prng.draw::<BaseElement>().unwrap(), constraint.cc());

    // single value constraints should evaluate to trace_value - value
    let trace_value = rand_value::<BaseElement>();
    assert_eq!(
        trace_value - value,
        constraint.evaluate_at(rand_value::<BaseElement>(), trace_value)
    );

    // constraint is build correctly for column 1 step 8
    let value = rand_value::<BaseElement>();
    let assertion = Assertion::single(1, 8, value);
    let constraint = BoundaryConstraint::<BaseElement, BaseElement>::new(
        assertion,
        inv_g,
        &mut twiddle_map,
        prng.draw().unwrap(),
    );
    assert_eq!(1, constraint.column());
    assert_eq!(vec![value], constraint.poly());
    assert_eq!((0, BaseElement::ONE), constraint.poly_offset());
    assert_eq!(&test_prng.draw::<BaseElement>().unwrap(), constraint.cc());

    // single value constraints should evaluate to trace_value - value
    let trace_value = rand_value::<BaseElement>();
    assert_eq!(
        trace_value - value,
        constraint.evaluate_at(rand_value::<BaseElement>(), trace_value)
    );

    // twiddle map was not touched
    assert!(twiddle_map.is_empty());
}

#[test]
fn boundary_constraint_from_periodic_assertion() {
    let mut test_prng = build_prng();
    let (inv_g, mut twiddle_map, mut prng) = build_constraint_params(16);

    // constraint should be built correctly for column 0, step 0, stride 4
    let value = rand_value::<BaseElement>();
    let assertion = Assertion::periodic(0, 0, 4, value);
    let constraint = BoundaryConstraint::<BaseElement, BaseElement>::new(
        assertion,
        inv_g,
        &mut twiddle_map,
        prng.draw().unwrap(),
    );
    assert_eq!(0, constraint.column());
    assert_eq!(vec![value], constraint.poly());
    assert_eq!((0, BaseElement::ONE), constraint.poly_offset());
    assert_eq!(&test_prng.draw::<BaseElement>().unwrap(), constraint.cc());

    // periodic value constraints should evaluate to trace_value - value
    let trace_value = rand_value::<BaseElement>();
    assert_eq!(
        trace_value - value,
        constraint.evaluate_at(rand_value::<BaseElement>(), trace_value)
    );

    // constraint should be built correctly for column 2, first step 3, stride 8
    let value = rand_value::<BaseElement>();
    let assertion = Assertion::periodic(2, 3, 8, value);
    let constraint = BoundaryConstraint::<BaseElement, BaseElement>::new(
        assertion,
        inv_g,
        &mut twiddle_map,
        prng.draw().unwrap(),
    );
    assert_eq!(2, constraint.column());
    assert_eq!(vec![value], constraint.poly());
    assert_eq!((0, BaseElement::ONE), constraint.poly_offset());
    assert_eq!(&test_prng.draw::<BaseElement>().unwrap(), constraint.cc());

    // periodic value constraints should evaluate to trace_value - value
    let trace_value = rand_value::<BaseElement>();
    assert_eq!(
        trace_value - value,
        constraint.evaluate_at(rand_value::<BaseElement>(), trace_value)
    );

    // twiddle map was not touched
    assert!(twiddle_map.is_empty());
}

#[test]
fn boundary_constraint_from_sequence_assertion() {
    let mut test_prng = build_prng();
    let (inv_g, mut twiddle_map, mut prng) = build_constraint_params(16);

    // constraint should be built correctly for column 0, first step 0, stride 4
    let values = rand_vector::<BaseElement>(4);
    let constraint_poly = build_sequence_poly(&values, 16);
    let assertion = Assertion::sequence(0, 0, 4, values);
    let constraint = BoundaryConstraint::<BaseElement, BaseElement>::new(
        assertion,
        inv_g,
        &mut twiddle_map,
        prng.draw().unwrap(),
    );
    assert_eq!(0, constraint.column());
    assert_eq!(constraint_poly, constraint.poly());
    assert_eq!((0, BaseElement::ONE), constraint.poly_offset());
    assert_eq!(&test_prng.draw::<BaseElement>().unwrap(), constraint.cc());
    assert_eq!(1, twiddle_map.len());

    // sequence value constraints with no offset should evaluate to
    // trace_value - constraint_poly(x)
    let x = rand_value::<BaseElement>();
    let trace_value = rand_value::<BaseElement>();
    assert_eq!(
        trace_value - polynom::eval(&constraint_poly, x),
        constraint.evaluate_at(x, trace_value)
    );

    // constraint should be built correctly for column 0, first step 3, stride 8
    let values = rand_vector::<BaseElement>(2);
    let constraint_poly = build_sequence_poly(&values, 16);
    let assertion = Assertion::sequence(0, 3, 8, values);
    let constraint = BoundaryConstraint::<BaseElement, BaseElement>::new(
        assertion,
        inv_g,
        &mut twiddle_map,
        prng.draw().unwrap(),
    );
    assert_eq!(0, constraint.column());
    assert_eq!(constraint_poly, constraint.poly());
    assert_eq!((3, inv_g.exp(3)), constraint.poly_offset());
    assert_eq!(&test_prng.draw::<BaseElement>().unwrap(), constraint.cc());
    assert_eq!(2, twiddle_map.len());

    // sequence value constraints with offset should evaluate to
    // trace_value - constraint_poly(x * offset)
    let x = rand_value::<BaseElement>();
    let trace_value = rand_value::<BaseElement>();
    assert_eq!(
        trace_value - polynom::eval(&constraint_poly, x * constraint.poly_offset().1),
        constraint.evaluate_at(x, trace_value)
    );
}

// PREPARE ASSERTIONS
// ================================================================================================

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

    let mut assertions = vec![
        Assertion::single(0, 9, BaseElement::new(5)), // column 0, step 9
        Assertion::single(0, 0, BaseElement::new(3)), // column 0, step 0
        Assertion::sequence(0, 3, 4, values.clone()), // column 1, steps 2, 6, 10, 14
        Assertion::sequence(0, 2, 4, values.clone()), // column 0, steps 2, 6, 10, 14
        Assertion::periodic(1, 3, 8, BaseElement::new(7)), // column 1, steps 3, 11
        Assertion::sequence(1, 0, 8, values[..2].to_vec()), // column 1, steps 0, 8
    ];

    // assertions should be sorted by stride, first step, and column
    let expected = vec![
        Assertion::single(0, 0, BaseElement::new(3)), // column 0, step 0
        Assertion::single(0, 9, BaseElement::new(5)), // column 0, step 9
        Assertion::sequence(0, 2, 4, values.clone()), // column 0, steps 2, 6, 10, 14
        Assertion::sequence(0, 3, 4, values.clone()), // column 1, steps 2, 6, 10, 14
        Assertion::sequence(1, 0, 8, values[..2].to_vec()), // column 1, steps 0, 8
        Assertion::periodic(1, 3, 8, BaseElement::new(7)), // column 1, steps 3, 11
    ];

    let trace_width = 2;
    let trace_length = 16;
    let result = super::prepare_assertions(assertions.clone(), trace_width, trace_length);
    assert_eq!(expected, result);

    shuffle(&mut assertions);
    let result = super::prepare_assertions(assertions.clone(), trace_width, trace_length);
    assert_eq!(expected, result);

    shuffle(&mut assertions);
    let result = super::prepare_assertions(assertions.clone(), trace_width, trace_length);
    assert_eq!(expected, result);
}

#[test]
#[should_panic(
    expected = "assertion (column=0, steps=[1, 9, ...], value=7) overlaps with assertion (column=0, step=9, value=5)"
)]
fn prepare_assertions_with_overlap() {
    let assertions = vec![
        Assertion::single(0, 9, BaseElement::new(5)),
        Assertion::periodic(0, 1, 8, BaseElement::new(7)),
    ];
    let _ = super::prepare_assertions(assertions, 2, 16);
}

#[test]
#[should_panic(
    expected = "assertion (column=0, step=16, value=5) is invalid: expected trace length to be at least 32, but was 16"
)]
fn prepare_assertions_with_invalid_trace_length() {
    let assertions = vec![Assertion::single(0, 16, BaseElement::new(5))];
    let _ = super::prepare_assertions(assertions, 2, 16);
}

#[test]
#[should_panic(
    expected = "assertion (column=3, step=17, value=5) is invalid: expected trace width to be at least 3, but was 2"
)]
fn prepare_assertions_with_invalid_trace_width() {
    let assertions = vec![Assertion::single(3, 17, BaseElement::new(5))];
    let _ = super::prepare_assertions(assertions, 2, 16);
}

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

#[allow(clippy::type_complexity)]
fn build_constraint_params(
    trace_length: usize,
) -> (
    BaseElement,
    BTreeMap<usize, Vec<BaseElement>>,
    DefaultRandomCoin<Blake3_256<BaseElement>>,
) {
    let inv_g = BaseElement::get_root_of_unity(trace_length.ilog2()).inv();
    let prng = build_prng();
    let twiddle_map = BTreeMap::<usize, Vec<BaseElement>>::new();
    (inv_g, twiddle_map, prng)
}


// 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::{super::super::ProofOptions, Vec, MIN_CYCLE_LENGTH};
use core::cmp;

// TRANSITION CONSTRAINT DEGREE
// ================================================================================================
/// Degree descriptor of a transition constraint.
///
/// Describes constraint degree as a combination of multiplications of periodic and trace
/// columns. For example, degree of a constraint which requires multiplication of two trace
/// columns can be described as: `base: 2, cycles: []`. A constraint which requires
/// multiplication of 3 trace columns and a periodic column with a period of 32 steps can be
/// described as: `base: 3, cycles: [32]`.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TransitionConstraintDegree {
    base: usize,
    cycles: Vec<usize>,
}

impl TransitionConstraintDegree {
    /// Creates a new transition constraint degree descriptor for constraints which involve
    /// multiplications of trace columns only.
    ///
    /// For example, if a constraint involves multiplication of two trace columns, `degree`
    /// should be set to 2. If a constraint involves multiplication of three trace columns,
    /// `degree` should be set to 3 etc.
    ///
    /// # Panics
    /// Panics if the provided `degree` is zero.
    pub fn new(degree: usize) -> Self {
        assert!(
            degree > 0,
            "transition constraint degree must be at least one, but was zero"
        );
        TransitionConstraintDegree {
            base: degree,
            cycles: vec![],
        }
    }

    /// Creates a new transition degree descriptor for constraints which involve multiplication
    /// of trace columns and periodic columns.
    ///
    /// For example, if a constraint involves multiplication of two trace columns and one
    /// periodic column with a period length of 32 steps, `base_degree` should be set to 2,
    /// and `cycles` should be set to `vec![32]`.
    ///
    /// # Panics
    /// Panics if:
    /// * `base_degree` is zero.
    /// * Any of the values in the `cycles` vector is smaller than two or is not powers of two.
    pub fn with_cycles(base_degree: usize, cycles: Vec<usize>) -> Self {
        assert!(
            base_degree > 0,
            "transition constraint degree must be at least one, but was zero"
        );
        for (i, &cycle) in cycles.iter().enumerate() {
            assert!(
                cycle >= MIN_CYCLE_LENGTH,
                "cycle length must be at least {MIN_CYCLE_LENGTH}, but was {cycle} for cycle {i}"
            );
            assert!(
                cycle.is_power_of_two(),
                "cycle length must be a power of two, but was {cycle} for cycle {i}"
            );
        }
        TransitionConstraintDegree {
            base: base_degree,
            cycles,
        }
    }

    /// Computes a degree to which this degree description expands in the context of execution
    /// trace of the specified length.
    ///
    /// The expanded degree is computed as follows:
    ///
    /// $$
    /// b \cdot (n - 1) + \sum_{i = 0}^{k - 1}{\frac{n \cdot (c_i - 1)}{c_i}}
    /// $$
    ///
    /// where: $b$ is the base degree, $n$ is the `trace_length`, $c_i$ is a cycle length of
    /// periodic column $i$, and $k$ is the total number of periodic columns for this degree
    /// descriptor.
    ///
    /// Thus, evaluation degree of a transition constraint which involves multiplication of two
    /// trace columns and one periodic column with a period length of 32 steps when evaluated
    /// over an execution trace of 64 steps would be:
    ///
    /// $$
    /// 2 \cdot (64 - 1) + \frac{64 \cdot (32 - 1)}{32} = 126 + 62 = 188
    /// $$
    pub fn get_evaluation_degree(&self, trace_length: usize) -> usize {
        let mut result = self.base * (trace_length - 1);
        for cycle_length in self.cycles.iter() {
            result += (trace_length / cycle_length) * (cycle_length - 1);
        }
        result
    }

    /// Returns a minimum blowup factor needed to evaluate constraint of this degree.
    ///
    /// This is guaranteed to be a power of two, greater than one.
    pub fn min_blowup_factor(&self) -> usize {
        // The blowup factor needs to be a power of two large enough to accommodate degree of
        // transition constraints defined by rational functions `C(x) / z(x)` where `C(x)` is the
        // constraint polynomial and `z(x)` is the transition constraint divisor.
        //
        // Degree of `C(x)` is always smaller than or equal to `[self.base + self.cycles.len()] * [trace_length - 1]`.
        // Degree of `z(x)` is `[trace_length - 1]`. Thus, the degree of `C(x) / z(x)` is
        // `[self.base + self.cycles.len() - 1] * [trace_length - 1]` and the blowup factor needed
        // to accommodate this degree can be estimated as `self.base + self.cycles.len() - 1`.
        //
        // For example, if degree of our constraints is 6, the blowup factor would need to be 8.
        // However, if the degree is 5, the blowup factor could be as small as 4.
        let degree_bound = self.base + self.cycles.len() - 1;
        cmp::max(
            degree_bound.next_power_of_two(),
            ProofOptions::MIN_BLOWUP_FACTOR,
        )
    }
}


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

// EVALUATION FRAME
// ================================================================================================
/// A set of execution trace rows required for evaluation of transition constraints.
///
/// In the current implementation, an evaluation frame always contains two consecutive rows of the
/// execution trace. It is passed in as one of the parameters into
/// [Air::evaluate_transition()](crate::Air::evaluate_transition) function.
#[derive(Debug, Clone)]
pub struct EvaluationFrame<E: FieldElement> {
    current: Vec<E>,
    next: Vec<E>,
}

impl<E: FieldElement> EvaluationFrame<E> {
    // CONSTRUCTORS
    // --------------------------------------------------------------------------------------------

    /// Returns a new evaluation frame instantiated with the specified number of columns.
    ///
    /// # Panics
    /// Panics if `num_columns` is zero.
    pub fn new(num_columns: usize) -> Self {
        assert!(
            num_columns > 0,
            "number of columns must be greater than zero"
        );
        EvaluationFrame {
            current: E::zeroed_vector(num_columns),
            next: E::zeroed_vector(num_columns),
        }
    }

    /// Returns a new evaluation frame instantiated from the provided rows.
    ///
    /// # Panics
    /// Panics if:
    /// * Lengths of the provided rows are zero.
    /// * Lengths of the provided rows are not the same.
    pub fn from_rows(current: Vec<E>, next: Vec<E>) -> Self {
        assert!(!current.is_empty(), "a row must contain at least one value");
        assert_eq!(
            current.len(),
            next.len(),
            "number of values in the rows must be the same"
        );
        Self { current, next }
    }

    // ROW ACCESSORS
    // --------------------------------------------------------------------------------------------

    /// Returns a reference to the current row.
    #[inline(always)]
    pub fn current(&self) -> &[E] {
        &self.current
    }

    /// Returns a reference to the next row.
    #[inline(always)]
    pub fn next(&self) -> &[E] {
        &self.next
    }

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

    /// Returns a mutable reference to the current row.
    #[inline(always)]
    pub fn current_mut(&mut self) -> &mut [E] {
        &mut self.current
    }

    /// Returns a mutable reference to the next row.
    #[inline(always)]
    pub fn next_mut(&mut self) -> &mut [E] {
        &mut self.next
    }
}


// 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::{AirContext, ConstraintDivisor, ExtensionOf, FieldElement, Vec};

mod frame;
pub use frame::EvaluationFrame;

mod degree;
pub use degree::TransitionConstraintDegree;

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

const MIN_CYCLE_LENGTH: usize = 2;

// TRANSITION CONSTRAINT INFO
// ================================================================================================
/// Metadata for transition constraints of a computation.
///
/// This metadata includes:
/// - List of transition constraint degrees for the main trace segment, as well as for auxiliary
///   trace segments (if any).
/// - Groupings of random composition constraint coefficients separately for the main trace segment
///   and for auxiliary tace segment.
/// - Divisor of transition constraints for a computation.
pub struct TransitionConstraints<E: FieldElement> {
    main_constraint_coef: Vec<E>,
    main_constraint_degrees: Vec<TransitionConstraintDegree>,
    aux_constraint_coef: Vec<E>,
    aux_constraint_degrees: Vec<TransitionConstraintDegree>,
    divisor: ConstraintDivisor<E::BaseField>,
}

impl<E: FieldElement> TransitionConstraints<E> {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new instance of [TransitionConstraints] for a computation described by the
    /// specified AIR context.
    ///
    /// # Panics
    /// Panics if the number of transition constraints in the context does not match the number of
    /// provided composition coefficients.
    pub fn new(context: &AirContext<E::BaseField>, composition_coefficients: &[E]) -> Self {
        assert_eq!(
            context.num_transition_constraints(),
            composition_coefficients.len(),
            "number of transition constraints must match the number of composition coefficient tuples"
        );

        // build constraint divisor; the same divisor applies to all transition constraints
        let divisor = ConstraintDivisor::from_transition(
            context.trace_len(),
            context.num_transition_exemptions(),
        );

        let main_constraint_degrees = context.main_transition_constraint_degrees.clone();
        let aux_constraint_degrees = context.aux_transition_constraint_degrees.clone();

        let (main_constraint_coef, aux_constraint_coef) =
            composition_coefficients.split_at(context.main_transition_constraint_degrees.len());
        Self {
            main_constraint_coef: main_constraint_coef.to_vec(),
            main_constraint_degrees,
            aux_constraint_coef: aux_constraint_coef.to_vec(),
            aux_constraint_degrees,
            divisor,
        }
    }

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

    /// Returns a list of transition constraint degree descriptors for the main trace segment of
    /// a computation.
    ///
    /// This list will be identical to the list passed into the [AirContext::new()] method as
    /// the `transition_constraint_degrees` parameter, or into [AirContext::new_multi_segment()]
    /// as the `main_transition_constraint_degrees` parameter.
    pub fn main_constraint_degrees(&self) -> &[TransitionConstraintDegree] {
        &self.main_constraint_degrees
    }

    /// Returns the number of constraints applied against the main trace segment of a computation.
    pub fn num_main_constraints(&self) -> usize {
        self.main_constraint_degrees.len()
    }

    /// Returns the random coefficients for constraints applied against main trace segment of a
    /// computation.
    pub fn main_constraint_coef(&self) -> Vec<E> {
        self.main_constraint_coef.clone()
    }

    /// Returns a list of transition constraint degree descriptors for auxiliary trace segments of
    /// a computation.
    ///
    /// This list will be identical to the list passed into [AirContext::new_multi_segment()]
    /// as the `aux_transition_constraint_degrees` parameter.
    pub fn aux_constraint_degrees(&self) -> &[TransitionConstraintDegree] {
        &self.aux_constraint_degrees
    }

    /// Returns the number of constraints applied against auxiliary trace segments of a
    /// computation.
    pub fn num_aux_constraints(&self) -> usize {
        self.aux_constraint_degrees.len()
    }

    /// Returns the random coefficients for constraints applied against auxiliary trace segments of a
    /// computation.
    pub fn aux_constraint_coef(&self) -> Vec<E> {
        self.aux_constraint_coef.clone()
    }

    /// Returns a divisor for transition constraints.
    ///
    /// All transition constraints have the same divisor which has the form:
    /// $$
    /// z(x) = \frac{x^n - 1}{x - g^{n - 1}}
    /// $$
    /// where: $n$ is the length of the execution trace and $g$ is the generator of the trace
    /// domain.
    ///
    /// This divisor specifies that transition constraints must hold on all steps of the
    /// execution trace except for the last one.
    pub fn divisor(&self) -> &ConstraintDivisor<E::BaseField> {
        &self.divisor
    }

    // CONSTRAINT COMPOSITION
    // --------------------------------------------------------------------------------------------

    /// Computes a linear combination of all transition constraint evaluations and divides the
    /// result by transition constraint divisor.
    ///
    /// A transition constraint is described by a rational function of the form $\frac{C(x)}{z(x)}$,
    /// where:
    /// * $C(x)$ is the constraint polynomial.
    /// * $z(x)$ is the constraint divisor polynomial.
    ///
    /// Thus, this function computes a linear combination of $C(x)$ evaluations.
    ///
    /// Since, the divisor polynomial is the same for all transition constraints (see
    /// [ConstraintDivisor::from_transition]), we can divide the linear combination by the
    /// divisor rather than dividing each individual $C(x)$ evaluation. This requires executing only
    /// one division at the end.
    pub fn combine_evaluations<F>(&self, main_evaluations: &[F], aux_evaluations: &[E], x: F) -> E
    where
        F: FieldElement<BaseField = E::BaseField>,
        E: ExtensionOf<F>,
    {
        // merge constraint evaluations for the main trace segment
        let mut result = main_evaluations
            .iter()
            .zip(self.main_constraint_coef.iter())
            .fold(E::ZERO, |acc, (&const_eval, &coef)| {
                acc + coef.mul_base(const_eval)
            });

        if !self.aux_constraint_coef.is_empty() {
            result += aux_evaluations
                .iter()
                .zip(self.aux_constraint_coef.iter())
                .fold(E::ZERO, |acc, (&const_eval, &coef)| acc + coef * const_eval);
        };
        // divide out the evaluation of divisor at x and return the result
        let z = E::from(self.divisor.evaluate_at(x));

        result / z
    }
}


// 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 crypto::Hasher;
use utils::{
    collections::Vec, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
    SliceReader,
};

// COMMITMENTS
// ================================================================================================
/// Commitments made by the prover during commit phase of the protocol.
///
/// These commitments include:
/// * Commitment to the extended execution trace, which may include commitments to one or more
///   execution trace segments.
/// * Commitment to the evaluations of constraint composition polynomial over LDE domain.
/// * Commitments to the evaluations of polynomials at all FRI layers.
///
/// Internally, the commitments are stored as a sequence of bytes. Thus, to retrieve the
/// commitments, [parse()](Commitments::parse) function should be used.
#[derive(Debug, Clone, Default, Eq, PartialEq)]
pub struct Commitments(Vec<u8>);

impl Commitments {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new Commitments struct initialized with the provided commitments.
    pub fn new<H: Hasher>(
        trace_roots: Vec<H::Digest>,
        constraint_root: H::Digest,
        fri_roots: Vec<H::Digest>,
    ) -> Self {
        let mut bytes = Vec::new();
        bytes.write(trace_roots);
        bytes.write(constraint_root);
        bytes.write(fri_roots);
        Commitments(bytes)
    }

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

    /// Adds the specified commitment to the list of commitments.
    pub fn add<H: Hasher>(&mut self, commitment: &H::Digest) {
        commitment.write_into(&mut self.0);
    }

    // PARSING
    // --------------------------------------------------------------------------------------------

    /// Parses the serialized commitments into distinct parts.
    ///
    /// The parts are (in the order in which they appear in the tuple):
    /// 1. Extended execution trace commitments.
    /// 2. Constraint composition polynomial evaluation commitment.
    /// 3. FRI layer commitments.
    ///
    /// # Errors
    /// Returns an error if the bytes stored in self could not be parsed into the requested number
    /// of commitments, or if there are any unconsumed bytes remaining after the parsing completes.
    #[allow(clippy::type_complexity)]
    pub fn parse<H: Hasher>(
        self,
        num_trace_segments: usize,
        num_fri_layers: usize,
    ) -> Result<(Vec<H::Digest>, H::Digest, Vec<H::Digest>), DeserializationError> {
        let mut reader = SliceReader::new(&self.0);

        // parse trace commitments
        let trace_commitments = H::Digest::read_batch_from(&mut reader, num_trace_segments)?;

        // parse constraint evaluation commitment:
        let constraint_commitment = H::Digest::read_from(&mut reader)?;

        // read FRI commitments (+ 1 for remainder polynomial commitment)
        let fri_commitments = H::Digest::read_batch_from(&mut reader, num_fri_layers + 1)?;

        // make sure we consumed all available commitment bytes
        if reader.has_more_bytes() {
            return Err(DeserializationError::UnconsumedBytes);
        }
        Ok((trace_commitments, constraint_commitment, fri_commitments))
    }
}

impl Serializable for Commitments {
    /// Serializes `self` and writes the resulting bytes into the `target`.
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        assert!(self.0.len() < u16::MAX as usize);
        target.write_u16(self.0.len() as u16);
        target.write_bytes(&self.0);
    }
}

impl Deserializable for Commitments {
    /// Reads commitments from the specified `source` and returns the result.
    ///
    /// # Errors
    /// Returns an error of a valid Commitments struct could not be read from the specified
    /// `source`.
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let num_bytes = source.read_u16()? as usize;
        let result = source.read_vec(num_bytes)?;
        Ok(Commitments(result))
    }
}


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

use crate::{ProofOptions, TraceInfo, TraceLayout};
use math::{StarkField, ToElements};
use utils::{
    collections::Vec, string::ToString, ByteReader, ByteWriter, Deserializable,
    DeserializationError, Serializable,
};

// PROOF CONTEXT
// ================================================================================================
/// Basic metadata about a specific execution of a computation.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Context {
    trace_layout: TraceLayout,
    trace_length: usize,
    trace_meta: Vec<u8>,
    field_modulus_bytes: Vec<u8>,
    options: ProofOptions,
}

impl Context {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Creates a new context for a computation described by the specified field, trace info, and
    /// proof options.
    pub fn new<B: StarkField>(trace_info: &TraceInfo, options: ProofOptions) -> Self {
        Context {
            trace_layout: trace_info.layout().clone(),
            trace_length: trace_info.length(),
            trace_meta: trace_info.meta().to_vec(),
            field_modulus_bytes: B::get_modulus_le_bytes(),
            options,
        }
    }

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

    /// Returns a layout describing how columns of the execution trace described by this context
    /// are arranged into segments.
    pub fn trace_layout(&self) -> &TraceLayout {
        &self.trace_layout
    }

    /// Returns execution trace length of the computation described by this context.
    pub fn trace_length(&self) -> usize {
        self.trace_length
    }

    /// Returns execution trace info for the computation described by this context.
    pub fn get_trace_info(&self) -> TraceInfo {
        TraceInfo::new_multi_segment(
            self.trace_layout.clone(),
            self.trace_length(),
            self.trace_meta.clone(),
        )
    }

    /// Returns the size of the LDE domain for the computation described by this context.
    pub fn lde_domain_size(&self) -> usize {
        self.trace_length() * self.options.blowup_factor()
    }

    /// Returns modulus of the field for the computation described by this context.
    pub fn field_modulus_bytes(&self) -> &[u8] {
        &self.field_modulus_bytes
    }

    /// Returns number of bits in the base field modulus for the computation described by this
    /// context.
    ///
    /// The modulus is assumed to be encoded in little-endian byte order.
    pub fn num_modulus_bits(&self) -> u32 {
        let mut num_bits = self.field_modulus_bytes.len() as u32 * 8;
        for &byte in self.field_modulus_bytes.iter().rev() {
            if byte != 0 {
                num_bits -= byte.leading_zeros();
                return num_bits;
            }
            num_bits -= 8;
        }

        0
    }

    /// Returns proof options which were used to a proof in this context.
    pub fn options(&self) -> &ProofOptions {
        &self.options
    }
}

impl<E: StarkField> ToElements<E> for Context {
    /// Converts this [Context] into a vector of field elements.
    ///
    /// The elements are layed out as follows:
    /// - trace layout info [1 or more elements].
    /// - field modulus bytes [2 field elements].
    /// - field extension and FRI parameters [1 element].
    /// - grinding factor [1 element].
    /// - blowup factor [1 element].
    /// - number of queries [1 element].
    /// - trace length [1 element].
    /// - trace metadata [0 or more elements].
    fn to_elements(&self) -> Vec<E> {
        // convert trace layout
        let mut result = self.trace_layout.to_elements();

        // convert field modulus bytes into 2 elements
        let num_modulus_bytes = self.field_modulus_bytes.len();
        let (m1, m2) = self.field_modulus_bytes.split_at(num_modulus_bytes / 2);
        result.push(bytes_to_element(m1));
        result.push(bytes_to_element(m2));

        // convert proof options and trace length to elements
        result.append(&mut self.options.to_elements());
        result.push(E::from(self.trace_length as u64));

        // convert trace metadata to elements; this is done by breaking trace metadata into chunks
        // of bytes which are slightly smaller than the number of bytes needed to encode a field
        // element, and then converting these chunks into field elements.
        if !self.trace_meta.is_empty() {
            for chunk in self.trace_meta.chunks(E::ELEMENT_BYTES - 1) {
                result.push(bytes_to_element(chunk));
            }
        }

        result
    }
}

impl Serializable for Context {
    /// Serializes `self` and writes the resulting bytes into the `target`.
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        self.trace_layout.write_into(target);
        target.write_u8(self.trace_length.ilog2() as u8); // store as power of two
        target.write_u16(self.trace_meta.len() as u16);
        target.write_bytes(&self.trace_meta);
        assert!(self.field_modulus_bytes.len() < u8::MAX as usize);
        target.write_u8(self.field_modulus_bytes.len() as u8);
        target.write_bytes(&self.field_modulus_bytes);
        self.options.write_into(target);
    }
}

impl Deserializable for Context {
    /// Reads proof context from the specified `source` and returns the result.
    ///
    /// # Errors
    /// Returns an error of a valid Context struct could not be read from the specified `source`.
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        // read and validate trace layout info
        let trace_layout = TraceLayout::read_from(source)?;

        // read and validate trace length (which was stored as a power of two)
        let trace_length = source.read_u8()?;
        if trace_length < TraceInfo::MIN_TRACE_LENGTH.ilog2() as u8 {
            return Err(DeserializationError::InvalidValue(format!(
                "trace length cannot be smaller than 2^{}, but was 2^{}",
                TraceInfo::MIN_TRACE_LENGTH.ilog2(),
                trace_length
            )));
        }
        let trace_length = 2_usize.pow(trace_length as u32);

        // read trace metadata
        let num_meta_bytes = source.read_u16()? as usize;
        let trace_meta = if num_meta_bytes != 0 {
            source.read_vec(num_meta_bytes)?
        } else {
            vec![]
        };

        // read and validate field modulus bytes
        let num_modulus_bytes = source.read_u8()? as usize;
        if num_modulus_bytes == 0 {
            return Err(DeserializationError::InvalidValue(
                "field modulus cannot be an empty value".to_string(),
            ));
        }
        let field_modulus_bytes = source.read_vec(num_modulus_bytes)?;

        // read options
        let options = ProofOptions::read_from(source)?;

        Ok(Context {
            trace_layout,
            trace_length,
            trace_meta,
            field_modulus_bytes,
            options,
        })
    }
}

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

/// Converts a slice of bytes into a field element.
///
/// Assumes that the length of `bytes` is smaller than the number of bytes needed to encode an
/// element.
#[allow(clippy::let_and_return)]
fn bytes_to_element<B: StarkField>(bytes: &[u8]) -> B {
    debug_assert!(bytes.len() < B::ELEMENT_BYTES);

    let mut buf = bytes.to_vec();
    buf.resize(B::ELEMENT_BYTES, 0);
    let element = match B::try_from(&buf) {
        Ok(element) => element,
        Err(_) => panic!("element deserialization failed"),
    };
    element
}

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

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

    #[test]
    fn context_to_elements() {
        let field_extension = FieldExtension::None;
        let fri_folding_factor = 8;
        let fri_remainder_max_degree = 127;
        let grinding_factor = 20;
        let blowup_factor = 8;
        let num_queries = 30;

        let main_width = 20;
        let num_aux_segments = 1;
        let aux_width = 9;
        let aux_rands = 12;
        let trace_length = 4096;

        let ext_fri = u32::from_le_bytes([
            fri_remainder_max_degree,
            fri_folding_factor,
            field_extension as u8,
            0,
        ]);

        let layout_info = u32::from_le_bytes([aux_rands, aux_width, num_aux_segments, main_width]);

        let expected = vec![
            BaseElement::from(layout_info),
            BaseElement::from(1_u32),    // lower bits of field modulus
            BaseElement::from(u32::MAX), // upper bits of field modulus
            BaseElement::from(ext_fri),
            BaseElement::from(grinding_factor as u32),
            BaseElement::from(blowup_factor as u32),
            BaseElement::from(num_queries as u32),
            BaseElement::from(trace_length as u32),
        ];

        let options = ProofOptions::new(
            num_queries,
            blowup_factor,
            grinding_factor,
            field_extension,
            fri_folding_factor as usize,
            fri_remainder_max_degree as usize,
        );
        let layout = TraceLayout::new(
            main_width as usize,
            [aux_width as usize],
            [aux_rands as usize],
        );
        let trace_info = TraceInfo::new_multi_segment(layout, trace_length, vec![]);
        let context = Context::new::<BaseElement>(&trace_info, options);
        assert_eq!(expected, context.to_elements());
    }
}


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

//! Contains STARK proof struct and associated components.

use crate::{ProofOptions, TraceInfo, TraceLayout};
use core::cmp;
use crypto::Hasher;
use fri::FriProof;
use utils::{
    collections::Vec, ByteReader, Deserializable, DeserializationError, Serializable, SliceReader,
};

mod context;
pub use context::Context;

mod commitments;
pub use commitments::Commitments;

mod queries;
pub use queries::Queries;

mod ood_frame;
pub use ood_frame::OodFrame;

mod table;
pub use table::Table;

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

const GRINDING_CONTRIBUTION_FLOOR: u32 = 80;

// STARK PROOF
// ================================================================================================
/// A proof generated by Winterfell prover.
///
/// A STARK proof contains information proving that a computation was executed correctly. A proof
/// also contains basic metadata for the computation, but neither the definition of the computation
/// itself, nor public inputs consumed by the computation are contained in a proof.
///
/// A proof can be serialized into a sequence of bytes using [to_bytes()](StarkProof::to_bytes)
/// function, and deserialized from a sequence of bytes using [from_bytes()](StarkProof::from_bytes)
/// function.
///
/// To estimate soundness of a proof (in bits), [security_level()](StarkProof::security_level)
/// function can be used.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct StarkProof {
    /// Basic metadata about the execution of the computation described by this proof.
    pub context: Context,
    /// Commitments made by the prover during the commit phase of the protocol.
    pub commitments: Commitments,
    /// Decommitments of extended execution trace values (for all trace segments) at position
    ///  queried by the verifier.
    pub trace_queries: Vec<Queries>,
    /// Decommitments of constraint composition polynomial evaluations at positions queried by
    /// the verifier.
    pub constraint_queries: Queries,
    /// Trace and constraint polynomial evaluations at an out-of-domain point.
    pub ood_frame: OodFrame,
    /// Low-degree proof for a DEEP composition polynomial.
    pub fri_proof: FriProof,
    /// Proof-of-work nonce for query seed grinding.
    pub pow_nonce: u64,
}

impl StarkProof {
    /// Returns STARK protocol parameters used to generate this proof.
    pub fn options(&self) -> &ProofOptions {
        self.context.options()
    }

    /// Returns a layout describing how columns of the execution trace described by this context
    /// are arranged into segments.
    pub fn trace_layout(&self) -> &TraceLayout {
        self.context.trace_layout()
    }

    /// Returns trace length for the computation described by this proof.
    pub fn trace_length(&self) -> usize {
        self.context.trace_length()
    }

    /// Returns trace info for the computation described by this proof.
    pub fn get_trace_info(&self) -> TraceInfo {
        self.context.get_trace_info()
    }

    /// Returns the size of the LDE domain for the computation described by this proof.
    pub fn lde_domain_size(&self) -> usize {
        self.context.lde_domain_size()
    }

    // SECURITY LEVEL
    // --------------------------------------------------------------------------------------------
    /// Returns security level of this proof (in bits).
    ///
    /// When `conjectured` is true, conjectured security level is returned; otherwise, provable
    /// security level is returned. Usually, the number of queries needed for provable security is
    /// 2x - 3x higher than the number of queries needed for conjectured security at the same
    /// security level.
    pub fn security_level<H: Hasher>(&self, conjectured: bool) -> u32 {
        if conjectured {
            get_conjectured_security(
                self.context.options(),
                self.context.num_modulus_bits(),
                self.trace_length() as u64,
                H::COLLISION_RESISTANCE,
            )
        } else {
            #[cfg(not(feature = "std"))]
            panic!("proven security level is not available in no_std mode");

            #[cfg(feature = "std")]
            get_proven_security(
                self.context.options(),
                self.context.num_modulus_bits(),
                self.lde_domain_size() as u64,
                self.trace_length() as u64,
                H::COLLISION_RESISTANCE,
            )
        }
    }

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

    /// Serializes this proof into a vector of bytes.
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut result = Vec::new();
        self.context.write_into(&mut result);
        self.commitments.write_into(&mut result);
        self.trace_queries.write_into(&mut result);
        self.constraint_queries.write_into(&mut result);
        self.ood_frame.write_into(&mut result);
        self.fri_proof.write_into(&mut result);
        result.extend_from_slice(&self.pow_nonce.to_le_bytes());
        result
    }

    /// Returns a STARK proof read from the specified `source`.
    ///
    /// # Errors
    /// Returns an error of a valid STARK proof could not be read from the specified `source`.
    pub fn from_bytes(source: &[u8]) -> Result<Self, DeserializationError> {
        let mut source = SliceReader::new(source);

        // parse the context
        let context = Context::read_from(&mut source)?;

        // parse the commitments
        let commitments = Commitments::read_from(&mut source)?;

        // parse trace queries
        let num_trace_segments = context.trace_layout().num_segments();
        let mut trace_queries = Vec::with_capacity(num_trace_segments);
        for _ in 0..num_trace_segments {
            trace_queries.push(Queries::read_from(&mut source)?);
        }

        // parse the rest of the proof
        let proof = StarkProof {
            context,
            commitments,
            trace_queries,
            constraint_queries: Queries::read_from(&mut source)?,
            ood_frame: OodFrame::read_from(&mut source)?,
            fri_proof: FriProof::read_from(&mut source)?,
            pow_nonce: source.read_u64()?,
        };
        if source.has_more_bytes() {
            return Err(DeserializationError::UnconsumedBytes);
        }
        Ok(proof)
    }
}

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

/// Computes conjectured security level for the specified proof parameters.
fn get_conjectured_security(
    options: &ProofOptions,
    base_field_bits: u32,
    trace_domain_size: u64,
    collision_resistance: u32,
) -> u32 {
    // compute max security we can get for a given field size
    let field_size = base_field_bits * options.field_extension().degree();
    let field_security = field_size - trace_domain_size.trailing_zeros();

    // compute security we get by executing multiple query rounds
    let security_per_query = options.blowup_factor().ilog2();
    let mut query_security = security_per_query * options.num_queries() as u32;

    // include grinding factor contributions only for proofs adequate security
    if query_security >= GRINDING_CONTRIBUTION_FLOOR {
        query_security += options.grinding_factor();
    }

    cmp::min(
        cmp::min(field_security, query_security) - 1,
        collision_resistance,
    )
}

#[cfg(feature = "std")]
/// Estimates proven security level for the specified proof parameters.
fn get_proven_security(
    options: &ProofOptions,
    base_field_bits: u32,
    lde_domain_size: u64,
    trace_domain_size: u64,
    collision_resistance: u32,
) -> u32 {
    let extension_field_bits = (base_field_bits * options.field_extension().degree()) as f64;

    let blowup_bits = options.blowup_factor().ilog2() as f64;
    let num_fri_queries = options.num_queries() as f64;
    let lde_size_bits = lde_domain_size.trailing_zeros() as f64;

    // blowup_plus_bits is the number of bits in the blowup factor which is the inverse of
    // `\rho^+ := (trace_domain_size + 2) / lde_domain_size`. `\rho^+` is used in order to define a larger
    // agreement parameter `\alpha^+ := (1 + 1/2m)\sqrt{rho^+} := 1 - \theta^+`. The reason for
    // running FRI with a larger agreement parameter is to account for the simplified
    // DEEP composition polynomial. See Protocol 3 in https://eprint.iacr.org/2022/1216.
    let blowup_plus_bits = ((lde_domain_size as f64) / (trace_domain_size as f64 + 2_f64)).log2();

    // m is a parameter greater or equal to 3.
    // A larger m gives a worse field security bound but a better query security bound.
    // An optimal value of m is then a value that would balance field and query security
    // but there is no simple closed form solution.
    // This sets m so that field security is equal to the best query security for any value
    // of m, unless the calculated value is less than 3 in which case it gets rounded up to 3.
    let mut m = extension_field_bits + 1.0;
    m -= options.grinding_factor() as f64;
    m -= 1.5 * blowup_bits;
    m -= 0.5 * num_fri_queries * blowup_plus_bits;
    m -= 2.0 * lde_size_bits;
    m /= 7.0;
    m = 2.0_f64.powf(m);
    m -= 0.5;
    m = m.max(3.0);

    // compute pre-FRI query security
    // this considers only the third component given in the corresponding part of eq. 20
    // in https://eprint.iacr.org/2021/582, i.e. (m+1/2)^7.n^2 / (2\rho^1.5.q) as all
    // other terms are negligible in comparison.
    let pre_query_security = (extension_field_bits + 1.0
        - 3.0 / 2.0 * blowup_bits
        - 2.0 * lde_size_bits
        - 7.0 * (m + 0.5).log2()) as u32;

    // compute security we get by executing multiple query rounds
    let security_per_query = 0.5 * blowup_plus_bits - (1.0 + 1.0 / (2.0 * m)).log2();
    let mut query_security = (security_per_query * num_fri_queries) as u32;

    query_security += options.grinding_factor();

    cmp::min(
        cmp::min(pre_query_security, query_security) - 1,
        collision_resistance,
    )
}