// 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 structopt::StructOpt;
use winterfell::{
    crypto::hashers::{GriffinJive64_256, Rp64_256, RpJive64_256},
    math::fields::f128::BaseElement,
    FieldExtension, ProofOptions, StarkProof, VerifierError,
};

pub mod fibonacci;
#[cfg(feature = "std")]
pub mod lamport;
#[cfg(feature = "std")]
pub mod merkle;
pub mod rescue;
#[cfg(feature = "std")]
pub mod rescue_raps;
pub mod utils;
pub mod vdf;

#[cfg(test)]
mod tests;

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

pub type Blake3_192 = winterfell::crypto::hashers::Blake3_192<BaseElement>;
pub type Blake3_256 = winterfell::crypto::hashers::Blake3_256<BaseElement>;
pub type Sha3_256 = winterfell::crypto::hashers::Sha3_256<BaseElement>;

pub trait Example {
    fn prove(&self) -> StarkProof;
    fn verify(&self, proof: StarkProof) -> Result<(), VerifierError>;
    fn verify_with_wrong_inputs(&self, proof: StarkProof) -> Result<(), VerifierError>;
}

// EXAMPLE OPTIONS
// ================================================================================================

#[derive(StructOpt, Debug)]
#[structopt(name = "winterfell", about = "Winterfell examples")]
pub struct ExampleOptions {
    #[structopt(subcommand)]
    pub example: ExampleType,

    /// Hash function used in the protocol
    #[structopt(short = "h", long = "hash_fn", default_value = "blake3_256")]
    hash_fn: String,

    /// Number of queries to include in a proof
    #[structopt(short = "q", long = "queries")]
    num_queries: Option<usize>,

    /// Blowup factor for low degree extension
    #[structopt(short = "b", long = "blowup")]
    blowup_factor: Option<usize>,

    /// Grinding factor for query seed
    #[structopt(short = "g", long = "grinding", default_value = "16")]
    grinding_factor: u32,

    /// Field extension degree for composition polynomial
    #[structopt(short = "e", long = "field_extension", default_value = "1")]
    field_extension: u32,

    /// Folding factor for FRI protocol
    #[structopt(short = "f", long = "folding", default_value = "8")]
    folding_factor: usize,
}

impl ExampleOptions {
    pub fn to_proof_options(&self, q: usize, b: usize) -> (ProofOptions, HashFunction) {
        let num_queries = self.num_queries.unwrap_or(q);
        let blowup_factor = self.blowup_factor.unwrap_or(b);
        let field_extension = match self.field_extension {
            1 => FieldExtension::None,
            2 => FieldExtension::Quadratic,
            3 => FieldExtension::Cubic,
            val => panic!("'{val}' is not a valid field extension option"),
        };

        let hash_fn = match self.hash_fn.as_str() {
            "blake3_192" => HashFunction::Blake3_192,
            "blake3_256" => HashFunction::Blake3_256,
            "sha3_256" => HashFunction::Sha3_256,
            "rp64_256" => HashFunction::Rp64_256,
            "rp_jive64_256" => HashFunction::RpJive64_256,
            "griffin_jive64_256" => HashFunction::GriffinJive64_256,
            val => panic!("'{val}' is not a valid hash function option"),
        };

        (
            ProofOptions::new(
                num_queries,
                blowup_factor,
                self.grinding_factor,
                field_extension,
                self.folding_factor,
                31,
            ),
            hash_fn,
        )
    }

    /// Returns security level of the input proof in bits.
    pub fn get_proof_security_level(&self, proof: &StarkProof, conjectured: bool) -> usize {
        let security_level = match self.hash_fn.as_str() {
            "blake3_192" => proof.security_level::<Blake3_192>(conjectured),
            "blake3_256" => proof.security_level::<Blake3_256>(conjectured),
            "sha3_256" => proof.security_level::<Sha3_256>(conjectured),
            "rp64_256" => proof.security_level::<Rp64_256>(conjectured),
            "rp_jive64_256" => proof.security_level::<RpJive64_256>(conjectured),
            "griffin_jive64_256" => proof.security_level::<GriffinJive64_256>(conjectured),
            val => panic!("'{val}' is not a valid hash function option"),
        };

        security_level as usize
    }
}

#[derive(StructOpt, Debug)]
//#[structopt(about = "available examples")]
pub enum ExampleType {
    /// Compute a Fibonacci sequence using trace table with 2 registers
    Fib {
        /// Length of Fibonacci sequence; must be a power of two
        #[structopt(short = "n", default_value = "1048576")]
        sequence_length: usize,
    },
    /// Compute a Fibonacci sequence using trace table with 8 registers
    Fib8 {
        /// Length of Fibonacci sequence; must be a power of two
        #[structopt(short = "n", default_value = "1048576")]
        sequence_length: usize,
    },
    /// Compute a multiplicative Fibonacci sequence using trace table with 2 registers
    Mulfib {
        /// Length of Fibonacci sequence; must be a power of two
        #[structopt(short = "n", default_value = "1048576")]
        sequence_length: usize,
    },
    /// Compute a multiplicative Fibonacci sequence using trace table with 8 registers
    Mulfib8 {
        /// Length of Fibonacci sequence; must be a power of two
        #[structopt(short = "n", default_value = "1048576")]
        sequence_length: usize,
    },
    /// Compute a Fibonacci sequence using trace table with 2 registers in `f64` field.
    FibSmall {
        /// Length of Fibonacci sequence; must be a power of two
        #[structopt(short = "n", default_value = "65536")]
        sequence_length: usize,
    },
    /// Execute a simple VDF function
    Vdf {
        /// Number of steps in the VDF function; must be a power of two
        #[structopt(short = "n", default_value = "1048576")]
        num_steps: usize,
    },
    /// Similar to the VDF example, but exempts an extra row from transition constraints.
    VdfExempt {
        /// Number of steps in the VDF function; must be one less than a power of two
        #[structopt(short = "n", default_value = "1048575")]
        num_steps: usize,
    },
    /// Compute a hash chain using Rescue hash function
    Rescue {
        /// Length of the hash chain; must be a power of two
        #[structopt(short = "n", default_value = "1024")]
        chain_length: usize,
    },
    /// Compute two hash chains absorbing sequences that are a permutation of each other
    #[cfg(feature = "std")]
    RescueRaps {
        /// Length of the hash chain; must be a power of two and at least 4
        #[structopt(short = "n", default_value = "1024")]
        chain_length: usize,
    },
    /// Compute a root of a Merkle path using Rescue hash function
    #[cfg(feature = "std")]
    Merkle {
        /// Depth of the Merkle tree; must be one less than a power of two
        #[structopt(short = "n", default_value = "7")]
        tree_depth: usize,
    },
    /// Compute an aggregate Lamport+ signature
    #[cfg(feature = "std")]
    LamportA {
        /// Number of signatures to aggregate; must be a power of two
        #[structopt(short = "n", default_value = "4")]
        num_signatures: usize,
    },
    /// Compute a threshold Lamport+ signature
    #[cfg(feature = "std")]
    LamportT {
        /// Number of signers; must be one less than a power of two
        #[structopt(short = "n", default_value = "3")]
        num_signers: usize,
    },
}

/// Defines a set of hash functions available for the provided examples. Some examples may not
/// support all listed hash functions.
///
/// Choice of a hash function has a direct impact on proof generation time, proof size, and proof
/// soundness. In general, sounds of the proof is bounded by the collision resistance of the hash
/// function used by the protocol.
#[repr(u8)]
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum HashFunction {
    /// BLAKE3 hash function with 192 bit output.
    ///
    /// When this function is used in the STARK protocol, proof security cannot exceed 96 bits.
    Blake3_192,

    /// BLAKE3 hash function with 256 bit output.
    ///
    /// When this function is used in the STARK protocol, proof security cannot exceed 128 bits.
    Blake3_256,

    /// SHA3 hash function with 256 bit output.
    ///
    /// When this function is used in the STARK protocol, proof security cannot exceed 128 bits.
    Sha3_256,

    /// Rescue Prime hash function with 256 bit output. It only works in `f64` field.
    ///
    /// When this function is used in the STARK protocol, proof security cannot exceed 128 bits.
    Rp64_256,

    /// Rescue Prime hash function with 256 bit output. It only works in `f64` field.
    /// This instance uses the Jive compression mode in Merkle trees.
    ///
    /// When this function is used in the STARK protocol, proof security cannot exceed 128 bits.
    RpJive64_256,

    /// Griffin hash function with 256 bit output. It only works in `f64` field.
    /// This instance uses the Jive compression mode in Merkle trees.
    ///
    /// When this function is used in the STARK protocol, proof security cannot exceed 128 bits.
    GriffinJive64_256,
}


// 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 log::debug;
use std::io::Write;
use std::time::Instant;
use structopt::StructOpt;
use winterfell::StarkProof;

use examples::{fibonacci, rescue, vdf, ExampleOptions, ExampleType};
#[cfg(feature = "std")]
use examples::{lamport, merkle, rescue_raps};

// EXAMPLE RUNNER
// ================================================================================================

fn main() {
    // configure logging
    env_logger::Builder::new()
        .format(|buf, record| writeln!(buf, "{}", record.args()))
        .filter_level(log::LevelFilter::Debug)
        .init();

    // read command-line args
    let options = ExampleOptions::from_args();

    debug!("============================================================");

    // instantiate and prepare the example
    let example = match options.example {
        ExampleType::Fib { sequence_length } => {
            fibonacci::fib2::get_example(&options, sequence_length)
        }
        ExampleType::Fib8 { sequence_length } => {
            fibonacci::fib8::get_example(&options, sequence_length)
        }
        ExampleType::Mulfib { sequence_length } => {
            fibonacci::mulfib2::get_example(&options, sequence_length)
        }
        ExampleType::Mulfib8 { sequence_length } => {
            fibonacci::mulfib8::get_example(&options, sequence_length)
        }
        ExampleType::FibSmall { sequence_length } => {
            fibonacci::fib_small::get_example(&options, sequence_length)
        }
        ExampleType::Vdf { num_steps } => vdf::regular::get_example(&options, num_steps),
        ExampleType::VdfExempt { num_steps } => vdf::exempt::get_example(&options, num_steps),
        ExampleType::Rescue { chain_length } => rescue::get_example(&options, chain_length),
        #[cfg(feature = "std")]
        ExampleType::RescueRaps { chain_length } => {
            rescue_raps::get_example(&options, chain_length)
        }
        #[cfg(feature = "std")]
        ExampleType::Merkle { tree_depth } => merkle::get_example(&options, tree_depth),
        #[cfg(feature = "std")]
        ExampleType::LamportA { num_signatures } => {
            lamport::aggregate::get_example(&options, num_signatures)
        }
        #[cfg(feature = "std")]
        ExampleType::LamportT { num_signers } => {
            lamport::threshold::get_example(&options, num_signers)
        }
    }
    .expect("The example failed to initialize.");

    // generate proof
    let now = Instant::now();
    let example = example.as_ref();
    let proof = example.prove();
    debug!(
        "---------------------\nProof generated in {} ms",
        now.elapsed().as_millis()
    );

    let proof_bytes = proof.to_bytes();
    debug!("Proof size: {:.1} KB", proof_bytes.len() as f64 / 1024f64);
    let conjectured_security_level = options.get_proof_security_level(&proof, true);

    #[cfg(feature = "std")]
    {
        let proven_security_level = options.get_proof_security_level(&proof, false);
        debug!(
            "Proof security: {} bits ({} proven)",
            conjectured_security_level, proven_security_level,
        );
    }

    #[cfg(not(feature = "std"))]
    debug!("Proof security: {} bits", conjectured_security_level);

    #[cfg(feature = "std")]
    debug!(
        "Proof hash: {}",
        hex::encode(blake3::hash(&proof_bytes).as_bytes())
    );

    // verify the proof
    debug!("---------------------");
    let parsed_proof = StarkProof::from_bytes(&proof_bytes).unwrap();
    assert_eq!(proof, parsed_proof);
    let now = Instant::now();
    match example.verify(proof) {
        Ok(_) => debug!(
            "Proof verified in {:.1} ms",
            now.elapsed().as_micros() as f64 / 1000f64
        ),
        Err(msg) => debug!("Failed to verify proof: {}", msg),
    }
    debug!("============================================================");
}


// 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::Example;

pub fn test_basic_proof_verification(e: Box<dyn Example>) {
    let proof = e.prove();
    assert!(e.verify(proof).is_ok());
}

pub fn test_basic_proof_verification_fail(e: Box<dyn Example>) {
    let proof = e.prove();
    let verified = e.verify_with_wrong_inputs(proof);
    assert!(verified.is_err());
}


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

pub mod fib2;
pub mod fib8;
pub mod fib_small;
pub mod mulfib2;
pub mod mulfib8;

mod utils;


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

use winterfell::math::{fields::f128::BaseElement, FieldElement};

pub fn compute_fib_term<E: FieldElement>(n: usize) -> E {
    let mut t0 = E::ONE;
    let mut t1 = E::ONE;

    for _ in 0..(n - 1) {
        t1 = t0 + t1;
        core::mem::swap(&mut t0, &mut t1);
    }

    t1
}

pub fn compute_mulfib_term(n: usize) -> BaseElement {
    let mut t0 = BaseElement::ONE;
    let mut t1 = BaseElement::new(2);

    for _ in 0..(n - 1) {
        t1 = t0 * t1;
        core::mem::swap(&mut t0, &mut t1);
    }

    t1
}

#[cfg(test)]
pub fn build_proof_options(use_extension_field: bool) -> winterfell::ProofOptions {
    use winterfell::{FieldExtension, ProofOptions};

    let extension = if use_extension_field {
        FieldExtension::Quadratic
    } else {
        FieldExtension::None
    };
    ProofOptions::new(28, 8, 0, extension, 4, 7)
}


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

use super::{BaseElement, FieldElement, ProofOptions, TRACE_WIDTH};
use crate::utils::are_equal;
use winterfell::{
    Air, AirContext, Assertion, EvaluationFrame, TraceInfo, TransitionConstraintDegree,
};

// FIBONACCI AIR
// ================================================================================================

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

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

    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    fn new(trace_info: TraceInfo, pub_inputs: Self::BaseField, options: ProofOptions) -> Self {
        let degrees = vec![
            TransitionConstraintDegree::new(1),
            TransitionConstraintDegree::new(1),
        ];
        assert_eq!(TRACE_WIDTH, trace_info.width());
        FibAir {
            context: AirContext::new(trace_info, degrees, 3, options),
            result: pub_inputs,
        }
    }

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

    fn evaluate_transition<E: FieldElement + From<Self::BaseField>>(
        &self,
        frame: &EvaluationFrame<E>,
        _periodic_values: &[E],
        result: &mut [E],
    ) {
        let current = frame.current();
        let next = frame.next();
        // expected state width is 2 field elements
        debug_assert_eq!(TRACE_WIDTH, current.len());
        debug_assert_eq!(TRACE_WIDTH, next.len());

        // constraints of Fibonacci sequence (2 terms per step):
        // s_{0, i+1} = s_{0, i} + s_{1, i}
        // s_{1, i+1} = s_{1, i} + s_{0, i+1}
        result[0] = are_equal(next[0], current[0] + current[1]);
        result[1] = are_equal(next[1], current[1] + next[0]);
    }

    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
        // a valid Fibonacci sequence should start with two ones and terminate with
        // the expected result
        let last_step = self.trace_length() - 1;
        vec![
            Assertion::single(0, 0, Self::BaseField::ONE),
            Assertion::single(1, 0, Self::BaseField::ONE),
            Assertion::single(1, last_step, self.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 super::utils::compute_fib_term;
use crate::{Blake3_192, Blake3_256, Example, ExampleOptions, HashFunction, Sha3_256};
use core::marker::PhantomData;
use log::debug;
use std::time::Instant;
use winterfell::{
    crypto::{DefaultRandomCoin, ElementHasher},
    math::{fields::f128::BaseElement, FieldElement},
    ProofOptions, Prover, StarkProof, Trace, TraceTable, VerifierError,
};

mod air;
use air::FibAir;

mod prover;
use prover::FibProver;

#[cfg(test)]
mod tests;

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

const TRACE_WIDTH: usize = 2;

// FIBONACCI EXAMPLE
// ================================================================================================

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

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

pub struct FibExample<H: ElementHasher> {
    options: ProofOptions,
    sequence_length: usize,
    result: BaseElement,
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> FibExample<H> {
    pub fn new(sequence_length: usize, options: ProofOptions) -> Self {
        assert!(
            sequence_length.is_power_of_two(),
            "sequence length must be a power of 2"
        );

        // compute Fibonacci sequence
        let now = Instant::now();
        let result = compute_fib_term(sequence_length);
        debug!(
            "Computed Fibonacci sequence up to {}th term in {} ms",
            sequence_length,
            now.elapsed().as_millis()
        );

        FibExample {
            options,
            sequence_length,
            result,
            _hasher: PhantomData,
        }
    }
}

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

impl<H: ElementHasher> Example for FibExample<H>
where
    H: ElementHasher<BaseField = BaseElement>,
{
    fn prove(&self) -> StarkProof {
        debug!(
            "Generating proof for computing Fibonacci sequence (2 terms per step) up to {}th term\n\
            ---------------------",
            self.sequence_length
        );

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

        // generate execution trace
        let now = Instant::now();
        let trace = prover.build_trace(self.sequence_length);

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

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

    fn verify(&self, proof: StarkProof) -> Result<(), VerifierError> {
        winterfell::verify::<FibAir, H, DefaultRandomCoin<H>>(proof, self.result)
    }

    fn verify_with_wrong_inputs(&self, proof: StarkProof) -> Result<(), VerifierError> {
        winterfell::verify::<FibAir, H, DefaultRandomCoin<H>>(proof, self.result + BaseElement::ONE)
    }
}


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

use super::{
    BaseElement, DefaultRandomCoin, ElementHasher, FibAir, FieldElement, PhantomData, ProofOptions,
    Prover, Trace, TraceTable, TRACE_WIDTH,
};

// FIBONACCI PROVER
// ================================================================================================

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

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

    /// Builds an execution trace for computing a Fibonacci sequence of the specified length such
    /// that each row advances the sequence by 2 terms.
    pub fn build_trace(&self, sequence_length: usize) -> TraceTable<BaseElement> {
        assert!(
            sequence_length.is_power_of_two(),
            "sequence length must be a power of 2"
        );

        let mut trace = TraceTable::new(TRACE_WIDTH, sequence_length / 2);
        trace.fill(
            |state| {
                state[0] = BaseElement::ONE;
                state[1] = BaseElement::ONE;
            },
            |_, state| {
                state[0] += state[1];
                state[1] += state[0];
            },
        );

        trace
    }
}

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

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

    fn options(&self) -> &ProofOptions {
        &self.options
    }
}


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

use super::{super::utils::build_proof_options, Blake3_256};

#[test]
fn fib2_test_basic_proof_verification() {
    let fib = Box::new(super::FibExample::<Blake3_256>::new(
        16,
        build_proof_options(false),
    ));
    crate::tests::test_basic_proof_verification(fib);
}

#[test]
fn fib2_test_basic_proof_verification_extension() {
    let fib = Box::new(super::FibExample::<Blake3_256>::new(
        16,
        build_proof_options(true),
    ));
    crate::tests::test_basic_proof_verification(fib);
}

#[test]
fn fib2_test_basic_proof_verification_fail() {
    let fib = Box::new(super::FibExample::<Blake3_256>::new(
        16,
        build_proof_options(false),
    ));
    crate::tests::test_basic_proof_verification_fail(fib);
}


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

use super::{BaseElement, FieldElement, TRACE_WIDTH};
use crate::utils::are_equal;
use winterfell::{
    Air, AirContext, Assertion, EvaluationFrame, ProofOptions, TraceInfo,
    TransitionConstraintDegree,
};

// FIBONACCI AIR
// ================================================================================================

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

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

    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    fn new(trace_info: TraceInfo, pub_inputs: Self::BaseField, options: ProofOptions) -> Self {
        let degrees = vec![
            TransitionConstraintDegree::new(1),
            TransitionConstraintDegree::new(1),
        ];
        assert_eq!(TRACE_WIDTH, trace_info.width());
        Fib8Air {
            context: AirContext::new(trace_info, degrees, 3, options),
            result: pub_inputs,
        }
    }

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

    fn evaluate_transition<E: FieldElement + From<Self::BaseField>>(
        &self,
        frame: &EvaluationFrame<E>,
        _periodic_values: &[E],
        result: &mut [E],
    ) {
        let current = frame.current();
        let next = frame.next();
        // expected state width is 2 field elements
        debug_assert_eq!(TRACE_WIDTH, current.len());
        debug_assert_eq!(TRACE_WIDTH, next.len());

        // constraints of Fibonacci sequence (2 registers, skipping over 8 terms):
        let n0 = current[0] + current[1];
        let n1 = current[1] + n0;
        let n2 = n0 + n1;
        let n3 = n1 + n2;
        let n4 = n2 + n3;
        let n5 = n3 + n4;
        let n6 = n4 + n5;
        let n7 = n5 + n6;

        result[0] = are_equal(next[0], n6);
        result[1] = are_equal(next[1], n7);
    }

    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
        // assert that the trace starts with 7th and 8th terms of Fibonacci sequence (the first
        // 6 terms are not recorded in the trace), and ends with the expected result
        let last_step = self.trace_length() - 1;
        vec![
            Assertion::single(0, 0, BaseElement::new(13)),
            Assertion::single(1, 0, BaseElement::new(21)),
            Assertion::single(1, last_step, self.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 super::utils::compute_fib_term;
use crate::{Blake3_192, Blake3_256, Example, ExampleOptions, HashFunction, Sha3_256};
use core::marker::PhantomData;
use log::debug;
use std::time::Instant;
use winterfell::{
    crypto::{DefaultRandomCoin, ElementHasher},
    math::{fields::f128::BaseElement, FieldElement},
    ProofOptions, Prover, StarkProof, Trace, TraceTable, VerifierError,
};

mod air;
use air::Fib8Air;

mod prover;
use prover::Fib8Prover;

#[cfg(test)]
mod tests;

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

const TRACE_WIDTH: usize = 2;

// FIBONACCI EXAMPLE
// ================================================================================================

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

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

pub struct Fib8Example<H: ElementHasher> {
    options: ProofOptions,
    sequence_length: usize,
    result: BaseElement,
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> Fib8Example<H> {
    pub fn new(sequence_length: usize, options: ProofOptions) -> Self {
        assert!(
            sequence_length.is_power_of_two(),
            "sequence length must be a power of 2"
        );

        // compute Fibonacci sequence
        let now = Instant::now();
        let result = compute_fib_term(sequence_length);
        debug!(
            "Computed Fibonacci sequence up to {}th term in {} ms",
            sequence_length,
            now.elapsed().as_millis()
        );

        Fib8Example {
            options,
            sequence_length,
            result,
            _hasher: PhantomData,
        }
    }
}

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

impl<H: ElementHasher> Example for Fib8Example<H>
where
    H: ElementHasher<BaseField = BaseElement>,
{
    fn prove(&self) -> StarkProof {
        debug!(
            "Generating proof for computing Fibonacci sequence (8 terms per step) up to {}th term\n\
            ---------------------",
            self.sequence_length
        );

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

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

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

    fn verify(&self, proof: StarkProof) -> Result<(), VerifierError> {
        winterfell::verify::<Fib8Air, H, DefaultRandomCoin<H>>(proof, self.result)
    }

    fn verify_with_wrong_inputs(&self, proof: StarkProof) -> Result<(), VerifierError> {
        winterfell::verify::<Fib8Air, H, DefaultRandomCoin<H>>(
            proof,
            self.result + BaseElement::ONE,
        )
    }
}


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

use super::{
    BaseElement, DefaultRandomCoin, ElementHasher, Fib8Air, FieldElement, PhantomData,
    ProofOptions, Prover, Trace, TraceTable,
};

// FIBONACCI PROVER
// ================================================================================================

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

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

    /// Builds an execution trace for computing a Fibonacci sequence of the specified length such
    /// that each row advances the sequence by 8 terms.
    pub fn build_trace(&self, length: usize) -> TraceTable<BaseElement> {
        assert!(
            length.is_power_of_two(),
            "sequence length must be a power of 2"
        );

        // initialize the trace with 7th and 8th terms of Fibonacci sequence (skipping the first 6)
        let n0 = BaseElement::ONE;
        let n1 = BaseElement::ONE;
        let n2 = n0 + n1;
        let n3 = n1 + n2;
        let n4 = n2 + n3;
        let n5 = n3 + n4;
        let n6 = n4 + n5;
        let n7 = n5 + n6;

        let mut reg0 = vec![n6];
        let mut reg1 = vec![n7];

        for i in 0..(length / 8 - 1) {
            let n0 = reg0[i] + reg1[i];
            let n1 = reg1[i] + n0;
            let n2 = n0 + n1;
            let n3 = n1 + n2;
            let n4 = n2 + n3;
            let n5 = n3 + n4;
            let n6 = n4 + n5;
            let n7 = n5 + n6;

            reg0.push(n6);
            reg1.push(n7);
        }

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

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

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

    fn options(&self) -> &ProofOptions {
        &self.options
    }
}


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

use super::{super::utils::build_proof_options, Blake3_256};

#[test]
fn fib8_test_basic_proof_verification() {
    let fib = Box::new(super::Fib8Example::<Blake3_256>::new(
        64,
        build_proof_options(false),
    ));
    crate::tests::test_basic_proof_verification(fib);
}

#[test]
fn fib8_test_basic_proof_verification_extension() {
    let fib = Box::new(super::Fib8Example::<Blake3_256>::new(
        64,
        build_proof_options(true),
    ));
    crate::tests::test_basic_proof_verification(fib);
}

#[test]
fn fib8_test_basic_proof_verification_fail() {
    let fib = Box::new(super::Fib8Example::<Blake3_256>::new(
        64,
        build_proof_options(false),
    ));
    crate::tests::test_basic_proof_verification_fail(fib);
}


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

use super::{BaseElement, FieldElement, ProofOptions, TRACE_WIDTH};
use crate::utils::are_equal;
use winterfell::{
    Air, AirContext, Assertion, EvaluationFrame, TraceInfo, TransitionConstraintDegree,
};

// FIBONACCI AIR
// ================================================================================================

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

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

    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    fn new(trace_info: TraceInfo, pub_inputs: Self::BaseField, options: ProofOptions) -> Self {
        let degrees = vec![
            TransitionConstraintDegree::new(1),
            TransitionConstraintDegree::new(1),
        ];
        assert_eq!(TRACE_WIDTH, trace_info.width());
        FibSmall {
            context: AirContext::new(trace_info, degrees, 3, options),
            result: pub_inputs,
        }
    }

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

    fn evaluate_transition<E: FieldElement + From<Self::BaseField>>(
        &self,
        frame: &EvaluationFrame<E>,
        _periodic_values: &[E],
        result: &mut [E],
    ) {
        let current = frame.current();
        let next = frame.next();
        // expected state width is 2 field elements
        debug_assert_eq!(TRACE_WIDTH, current.len());
        debug_assert_eq!(TRACE_WIDTH, next.len());

        // constraints of Fibonacci sequence (2 terms per step):
        // s_{0, i+1} = s_{0, i} + s_{1, i}
        // s_{1, i+1} = s_{1, i} + s_{0, i+1}
        result[0] = are_equal(next[0], current[0] + current[1]);
        result[1] = are_equal(next[1], current[1] + next[0]);
    }

    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
        // a valid Fibonacci sequence should start with two ones and terminate with
        // the expected result
        let last_step = self.trace_length() - 1;
        vec![
            Assertion::single(0, 0, Self::BaseField::ONE),
            Assertion::single(1, 0, Self::BaseField::ONE),
            Assertion::single(1, last_step, self.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 super::utils::compute_fib_term;
use crate::{Example, ExampleOptions, HashFunction};
use core::marker::PhantomData;
use log::debug;
use std::time::Instant;
use winterfell::{
    crypto::{DefaultRandomCoin, ElementHasher},
    math::{fields::f64::BaseElement, FieldElement},
    ProofOptions, Prover, StarkProof, Trace, TraceTable, VerifierError,
};

mod air;
use air::FibSmall;

mod prover;
use prover::FibSmallProver;

#[cfg(test)]
mod tests;

// CONSTANTS AND TYPES
// ================================================================================================

const TRACE_WIDTH: usize = 2;

type Blake3_192 = winterfell::crypto::hashers::Blake3_192<BaseElement>;
type Blake3_256 = winterfell::crypto::hashers::Blake3_256<BaseElement>;
type Sha3_256 = winterfell::crypto::hashers::Sha3_256<BaseElement>;
type Rp64_256 = winterfell::crypto::hashers::Rp64_256;
type RpJive64_256 = winterfell::crypto::hashers::RpJive64_256;
type GriffinJive64_256 = winterfell::crypto::hashers::GriffinJive64_256;

// FIBONACCI EXAMPLE
// ================================================================================================

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

    match hash_fn {
        HashFunction::Blake3_192 => Ok(Box::new(FibExample::<Blake3_192>::new(
            sequence_length,
            options,
        ))),
        HashFunction::Blake3_256 => Ok(Box::new(FibExample::<Blake3_256>::new(
            sequence_length,
            options,
        ))),
        HashFunction::Sha3_256 => Ok(Box::new(FibExample::<Sha3_256>::new(
            sequence_length,
            options,
        ))),
        HashFunction::Rp64_256 => Ok(Box::new(FibExample::<Rp64_256>::new(
            sequence_length,
            options,
        ))),
        HashFunction::RpJive64_256 => Ok(Box::new(FibExample::<RpJive64_256>::new(
            sequence_length,
            options,
        ))),
        HashFunction::GriffinJive64_256 => Ok(Box::new(FibExample::<GriffinJive64_256>::new(
            sequence_length,
            options,
        ))),
    }
}

pub struct FibExample<H: ElementHasher> {
    options: ProofOptions,
    sequence_length: usize,
    result: BaseElement,
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> FibExample<H> {
    pub fn new(sequence_length: usize, options: ProofOptions) -> Self {
        assert!(
            sequence_length.is_power_of_two(),
            "sequence length must be a power of 2"
        );

        // compute Fibonacci sequence
        let now = Instant::now();
        let result = compute_fib_term::<BaseElement>(sequence_length);
        debug!(
            "Computed Fibonacci sequence up to {}th term in {} ms",
            sequence_length,
            now.elapsed().as_millis()
        );

        FibExample {
            options,
            sequence_length,
            result,
            _hasher: PhantomData,
        }
    }
}

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

impl<H: ElementHasher> Example for FibExample<H>
where
    H: ElementHasher<BaseField = BaseElement>,
{
    fn prove(&self) -> StarkProof {
        debug!(
            "Generating proof for computing Fibonacci sequence (2 terms per step) up to {}th term\n\
            ---------------------",
            self.sequence_length
        );

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

        // generate execution trace
        let now = Instant::now();
        let trace = prover.build_trace(self.sequence_length);

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

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

    fn verify(&self, proof: StarkProof) -> Result<(), VerifierError> {
        winterfell::verify::<FibSmall, H, DefaultRandomCoin<H>>(proof, self.result)
    }

    fn verify_with_wrong_inputs(&self, proof: StarkProof) -> Result<(), VerifierError> {
        winterfell::verify::<FibSmall, H, DefaultRandomCoin<H>>(
            proof,
            self.result + BaseElement::ONE,
        )
    }
}


// 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::{
    air::FibSmall, BaseElement, DefaultRandomCoin, ElementHasher, FieldElement, PhantomData,
    ProofOptions, Prover, Trace, TraceTable, TRACE_WIDTH,
};

// FIBONACCI PROVER
// ================================================================================================

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

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

    /// Builds an execution trace for computing a Fibonacci sequence of the specified length such
    /// that each row advances the sequence by 2 terms.
    pub fn build_trace(&self, sequence_length: usize) -> TraceTable<BaseElement> {
        assert!(
            sequence_length.is_power_of_two(),
            "sequence length must be a power of 2"
        );

        let mut trace = TraceTable::new(TRACE_WIDTH, sequence_length / 2);
        trace.fill(
            |state| {
                state[0] = BaseElement::ONE;
                state[1] = BaseElement::ONE;
            },
            |_, state| {
                state[0] += state[1];
                state[1] += state[0];
            },
        );

        trace
    }
}

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

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

    fn options(&self) -> &ProofOptions {
        &self.options
    }
}


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

use super::{super::utils::build_proof_options, Rp64_256};

#[test]
fn fib_small_test_basic_proof_verification() {
    let fib = Box::new(super::FibExample::<Rp64_256>::new(
        128,
        build_proof_options(false),
    ));
    crate::tests::test_basic_proof_verification(fib);
}

#[test]
fn fib_small_test_basic_proof_verification_extension() {
    let fib = Box::new(super::FibExample::<Rp64_256>::new(
        128,
        build_proof_options(true),
    ));
    crate::tests::test_basic_proof_verification(fib);
}

#[test]
fn fib_small_test_basic_proof_verification_fail() {
    let fib = Box::new(super::FibExample::<Rp64_256>::new(
        128,
        build_proof_options(false),
    ));
    crate::tests::test_basic_proof_verification_fail(fib);
}


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

use crate::utils::are_equal;
use winterfell::{
    math::{fields::f128::BaseElement, FieldElement},
    Air, AirContext, Assertion, EvaluationFrame, ProofOptions, TraceInfo,
    TransitionConstraintDegree,
};

// FIBONACCI AIR
// ================================================================================================

const TRACE_WIDTH: usize = 2;

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

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

    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    fn new(trace_info: TraceInfo, pub_inputs: Self::BaseField, options: ProofOptions) -> Self {
        let degrees = vec![
            TransitionConstraintDegree::new(2),
            TransitionConstraintDegree::new(2),
        ];
        assert_eq!(TRACE_WIDTH, trace_info.width());
        MulFib2Air {
            context: AirContext::new(trace_info, degrees, 3, options),
            result: pub_inputs,
        }
    }

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

    fn evaluate_transition<E: FieldElement + From<Self::BaseField>>(
        &self,
        frame: &EvaluationFrame<E>,
        _periodic_values: &[E],
        result: &mut [E],
    ) {
        let current = frame.current();
        let next = frame.next();
        // expected state width is 2 field elements
        debug_assert_eq!(TRACE_WIDTH, current.len());
        debug_assert_eq!(TRACE_WIDTH, next.len());

        // constraints of multiplicative Fibonacci (with 2 registers) which state that:
        // s_{0, i+1} = s_{0, i} * s_{1, i}
        // s_{1, i+1} = s_{1, i} * s_{0, i+1}
        result[0] = are_equal(next[0], current[0] * current[1]);
        result[1] = are_equal(next[1], current[1] * next[0]);
    }

    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
        // a valid multiplicative Fibonacci sequence should start with 1, 2 and terminate
        // with the expected result
        let last_step = self.trace_length() - 1;
        vec![
            Assertion::single(0, 0, Self::BaseField::new(1)),
            Assertion::single(1, 0, Self::BaseField::new(2)),
            Assertion::single(0, last_step, self.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 super::utils::compute_mulfib_term;
use crate::{Blake3_192, Blake3_256, Example, ExampleOptions, HashFunction, Sha3_256};
use core::marker::PhantomData;
use log::debug;
use std::time::Instant;
use winterfell::{
    crypto::{DefaultRandomCoin, ElementHasher},
    math::{fields::f128::BaseElement, FieldElement},
    ProofOptions, Prover, StarkProof, Trace, TraceTable, VerifierError,
};

mod air;
use air::MulFib2Air;

mod prover;
use prover::MulFib2Prover;

#[cfg(test)]
mod tests;

// FIBONACCI EXAMPLE
// ================================================================================================

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

    match hash_fn {
        HashFunction::Blake3_192 => Ok(Box::new(MulFib2Example::<Blake3_192>::new(
            sequence_length,
            options,
        ))),
        HashFunction::Blake3_256 => Ok(Box::new(MulFib2Example::<Blake3_256>::new(
            sequence_length,
            options,
        ))),
        HashFunction::Sha3_256 => Ok(Box::new(MulFib2Example::<Sha3_256>::new(
            sequence_length,
            options,
        ))),
        _ => Err("The specified hash function cannot be used with this example.".to_string()),
    }
}
pub struct MulFib2Example<H: ElementHasher> {
    options: ProofOptions,
    sequence_length: usize,
    result: BaseElement,
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> MulFib2Example<H> {
    pub fn new(sequence_length: usize, options: ProofOptions) -> Self {
        assert!(
            sequence_length.is_power_of_two(),
            "sequence length must be a power of 2"
        );

        // compute Fibonacci sequence
        let now = Instant::now();
        let result = compute_mulfib_term(sequence_length);
        debug!(
            "Computed multiplicative Fibonacci sequence up to {}th term in {} ms",
            sequence_length,
            now.elapsed().as_millis()
        );

        MulFib2Example {
            options,
            sequence_length,
            result,
            _hasher: PhantomData,
        }
    }
}

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

impl<H: ElementHasher> Example for MulFib2Example<H>
where
    H: ElementHasher<BaseField = BaseElement>,
{
    fn prove(&self) -> StarkProof {
        let sequence_length = self.sequence_length;
        debug!(
            "Generating proof for computing multiplicative Fibonacci sequence (8 terms per step) up to {}th term\n\
            ---------------------",
            sequence_length
        );

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

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

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

    fn verify(&self, proof: StarkProof) -> Result<(), VerifierError> {
        winterfell::verify::<MulFib2Air, H, DefaultRandomCoin<H>>(proof, self.result)
    }

    fn verify_with_wrong_inputs(&self, proof: StarkProof) -> Result<(), VerifierError> {
        winterfell::verify::<MulFib2Air, H, DefaultRandomCoin<H>>(
            proof,
            self.result + BaseElement::ONE,
        )
    }
}


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

use super::{
    BaseElement, DefaultRandomCoin, ElementHasher, MulFib2Air, PhantomData, ProofOptions, Prover,
    Trace, TraceTable,
};

// FIBONACCI PROVER
// ================================================================================================

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

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

    /// Builds an execution trace for computing a multiplicative version of a Fibonacci sequence of
    /// the specified length such that each row advances the sequence by 2 terms.
    pub fn build_trace(&self, length: usize) -> TraceTable<BaseElement> {
        assert!(
            length.is_power_of_two(),
            "sequence length must be a power of 2"
        );

        let mut reg0 = vec![BaseElement::new(1)];
        let mut reg1 = vec![BaseElement::new(2)];

        for i in 0..(length / 2 - 1) {
            reg0.push(reg0[i] * reg1[i]);
            reg1.push(reg1[i] * reg0[i + 1]);
        }

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

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

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

    fn options(&self) -> &ProofOptions {
        &self.options
    }
}


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

use super::{super::utils::build_proof_options, Blake3_256};

#[test]
fn mulfib2_test_basic_proof_verification() {
    let fib = Box::new(super::MulFib2Example::<Blake3_256>::new(
        16,
        build_proof_options(false),
    ));
    crate::tests::test_basic_proof_verification(fib);
}

#[test]
fn mulfib2_test_basic_proof_verification_extension() {
    let fib = Box::new(super::MulFib2Example::<Blake3_256>::new(
        16,
        build_proof_options(true),
    ));
    crate::tests::test_basic_proof_verification(fib);
}

#[test]
fn mulfib2_test_basic_proof_verification_fail() {
    let fib = Box::new(super::MulFib2Example::<Blake3_256>::new(
        16,
        build_proof_options(false),
    ));
    crate::tests::test_basic_proof_verification_fail(fib);
}


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

use crate::utils::are_equal;
use winterfell::{
    math::{fields::f128::BaseElement, FieldElement},
    Air, AirContext, Assertion, EvaluationFrame, ProofOptions, TraceInfo,
    TransitionConstraintDegree,
};

// FIBONACCI AIR
// ================================================================================================

const TRACE_WIDTH: usize = 8;

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

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

    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    fn new(trace_info: TraceInfo, pub_inputs: Self::BaseField, options: ProofOptions) -> Self {
        let degrees = vec![
            TransitionConstraintDegree::new(2),
            TransitionConstraintDegree::new(2),
            TransitionConstraintDegree::new(2),
            TransitionConstraintDegree::new(2),
            TransitionConstraintDegree::new(2),
            TransitionConstraintDegree::new(2),
            TransitionConstraintDegree::new(2),
            TransitionConstraintDegree::new(2),
        ];
        assert_eq!(TRACE_WIDTH, trace_info.width());
        MulFib8Air {
            context: AirContext::new(trace_info, degrees, 3, options),
            result: pub_inputs,
        }
    }

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

    fn evaluate_transition<E: FieldElement + From<Self::BaseField>>(
        &self,
        frame: &EvaluationFrame<E>,
        _periodic_values: &[E],
        result: &mut [E],
    ) {
        let current = frame.current();
        let next = frame.next();
        // expected state width is 2 field elements
        debug_assert_eq!(TRACE_WIDTH, current.len());
        debug_assert_eq!(TRACE_WIDTH, next.len());

        // constraints of multiplicative Fibonacci (with 8 registers) which state that:
        // s_{0, i+1} = s_{6, i} * s_{7, i}
        // s_{1, i+1} = s_{7, i} * s_{0, i+1}
        // s_{2, i+1} = s_{0, i+1} * s_{1, i+1}
        // s_{3, i+1} = s_{1, i+1} * s_{2, i+1}
        // s_{4, i+1} = s_{2, i+1} * s_{3, i+1}
        // s_{5, i+1} = s_{3, i+1} * s_{4, i+1}
        // s_{6, i+1} = s_{4, i+1} * s_{5, i+1}
        // s_{7, i+1} = s_{5, i+1} * s_{6, i+1}
        result[0] = are_equal(next[0], current[6] * current[7]);
        result[1] = are_equal(next[1], current[7] * next[0]);
        result[2] = are_equal(next[2], next[0] * next[1]);
        result[3] = are_equal(next[3], next[1] * next[2]);
        result[4] = are_equal(next[4], next[2] * next[3]);
        result[5] = are_equal(next[5], next[3] * next[4]);
        result[6] = are_equal(next[6], next[4] * next[5]);
        result[7] = are_equal(next[7], next[5] * next[6]);
    }

    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
        // a valid multiplicative Fibonacci sequence should start with 1, 2 and terminate
        // with the expected result
        let last_step = self.trace_length() - 1;
        vec![
            Assertion::single(0, 0, BaseElement::new(1)),
            Assertion::single(1, 0, BaseElement::new(2)),
            Assertion::single(6, last_step, self.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 super::utils::compute_mulfib_term;
use crate::{Blake3_192, Blake3_256, Example, ExampleOptions, HashFunction, Sha3_256};
use core::marker::PhantomData;
use log::debug;
use std::time::Instant;
use winterfell::{
    crypto::{DefaultRandomCoin, ElementHasher},
    math::{fields::f128::BaseElement, FieldElement},
    ProofOptions, Prover, StarkProof, Trace, TraceTable, VerifierError,
};

mod air;
use air::MulFib8Air;

mod prover;
use prover::MulFib8Prover;

#[cfg(test)]
mod tests;

// FIBONACCI EXAMPLE
// ================================================================================================

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

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

pub struct MulFib8Example<H: ElementHasher> {
    options: ProofOptions,
    sequence_length: usize,
    result: BaseElement,
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> MulFib8Example<H> {
    pub fn new(sequence_length: usize, options: ProofOptions) -> Self {
        assert!(
            sequence_length.is_power_of_two(),
            "sequence length must be a power of 2"
        );

        // compute Fibonacci sequence
        let now = Instant::now();
        let result = compute_mulfib_term(sequence_length);
        debug!(
            "Computed multiplicative Fibonacci sequence up to {}th term in {} ms",
            sequence_length,
            now.elapsed().as_millis()
        );

        MulFib8Example {
            options,
            sequence_length,
            result,
            _hasher: PhantomData,
        }
    }
}

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

impl<H: ElementHasher> Example for MulFib8Example<H>
where
    H: ElementHasher<BaseField = BaseElement>,
{
    fn prove(&self) -> StarkProof {
        let sequence_length = self.sequence_length;
        debug!(
            "Generating proof for computing multiplicative Fibonacci sequence (2 terms per step) up to {}th term\n\
            ---------------------",
            sequence_length
        );

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

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

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

    fn verify(&self, proof: StarkProof) -> Result<(), VerifierError> {
        winterfell::verify::<MulFib8Air, H, DefaultRandomCoin<H>>(proof, self.result)
    }

    fn verify_with_wrong_inputs(&self, proof: StarkProof) -> Result<(), VerifierError> {
        winterfell::verify::<MulFib8Air, H, DefaultRandomCoin<H>>(
            proof,
            self.result + BaseElement::ONE,
        )
    }
}


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

use super::{
    BaseElement, DefaultRandomCoin, ElementHasher, MulFib8Air, PhantomData, ProofOptions, Prover,
    Trace, TraceTable,
};

// FIBONACCI PROVER
// ================================================================================================

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

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

    /// Builds an execution trace for computing a multiplicative version of a Fibonacci sequence of
    /// the specified length such that each row advances the sequence by 8 terms.
    pub fn build_trace(&self, length: usize) -> TraceTable<BaseElement> {
        assert!(
            length.is_power_of_two(),
            "sequence length must be a power of 2"
        );

        let mut reg0 = vec![BaseElement::new(1)];
        let mut reg1 = vec![BaseElement::new(2)];
        let mut reg2 = vec![reg0[0] * reg1[0]];
        let mut reg3 = vec![reg1[0] * reg2[0]];
        let mut reg4 = vec![reg2[0] * reg3[0]];
        let mut reg5 = vec![reg3[0] * reg4[0]];
        let mut reg6 = vec![reg4[0] * reg5[0]];
        let mut reg7 = vec![reg5[0] * reg6[0]];

        for i in 0..(length / 8 - 1) {
            reg0.push(reg6[i] * reg7[i]);
            reg1.push(reg7[i] * reg0[i + 1]);
            reg2.push(reg0[i + 1] * reg1[i + 1]);
            reg3.push(reg1[i + 1] * reg2[i + 1]);
            reg4.push(reg2[i + 1] * reg3[i + 1]);
            reg5.push(reg3[i + 1] * reg4[i + 1]);
            reg6.push(reg4[i + 1] * reg5[i + 1]);
            reg7.push(reg5[i + 1] * reg6[i + 1]);
        }

        TraceTable::init(vec![reg0, reg1, reg2, reg3, reg4, reg5, reg6, reg7])
    }
}

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

    fn get_pub_inputs(&self, trace: &Self::Trace) -> BaseElement {
        let last_step = trace.length() - 1;
        trace.get(6, last_step)
    }

    fn options(&self) -> &ProofOptions {
        &self.options
    }
}


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

use super::{super::utils::build_proof_options, Blake3_256};

#[test]
fn mulfib8_test_basic_proof_verification() {
    let fib = Box::new(super::MulFib8Example::<Blake3_256>::new(
        64,
        build_proof_options(false),
    ));
    crate::tests::test_basic_proof_verification(fib);
}

#[test]
fn mulfib8_test_basic_proof_verification_extension() {
    let fib = Box::new(super::MulFib8Example::<Blake3_256>::new(
        64,
        build_proof_options(true),
    ));
    crate::tests::test_basic_proof_verification(fib);
}

#[test]
fn mulfib8_test_basic_proof_verification_fail() {
    let fib = Box::new(super::MulFib8Example::<Blake3_256>::new(
        64,
        build_proof_options(false),
    ));
    crate::tests::test_basic_proof_verification_fail(fib);
}


// 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::Example;

pub mod aggregate;
pub mod threshold;

mod signature;
use signature::{message_to_elements, PrivateKey, Signature};

use crate::utils::rescue::{self, CYCLE_LENGTH, NUM_ROUNDS as NUM_HASH_ROUNDS};


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

use super::rescue::Rescue128;
use rand_utils::prng_vector;
use std::{cmp::Ordering, convert::TryInto};
use winterfell::{
    math::{fields::f128::BaseElement, FieldElement, StarkField},
    Serializable,
};

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

const MESSAGE_BITS: usize = 254;

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

type KeyData = [BaseElement; 2];

pub struct PrivateKey {
    sec_keys: Vec<KeyData>,
    pub_keys: Vec<KeyData>,
    pub_key_hash: PublicKey,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PublicKey(KeyData);

pub struct Signature {
    pub ones: Vec<KeyData>,
    pub zeros: Vec<KeyData>,
}

// PRIVATE KEY IMPLEMENTATION
// ================================================================================================

impl PrivateKey {
    /// Returns a private key generated from the specified `seed`.
    pub fn from_seed(seed: [u8; 32]) -> Self {
        let keys_elements: Vec<BaseElement> = prng_vector(seed, MESSAGE_BITS * 2);
        let mut sec_keys = Vec::with_capacity(MESSAGE_BITS);
        let mut pub_keys = Vec::with_capacity(MESSAGE_BITS);

        for i in (0..keys_elements.len()).step_by(2) {
            let sk = [keys_elements[i], keys_elements[i + 1]];
            sec_keys.push(sk);

            let pk = Rescue128::digest(&sk).to_elements();
            pub_keys.push(pk);
        }

        let pub_key_hash = hash_pub_keys(&pub_keys);

        PrivateKey {
            sec_keys,
            pub_keys,
            pub_key_hash,
        }
    }

    /// Returns a public key corresponding to this private key.
    pub fn pub_key(&self) -> PublicKey {
        self.pub_key_hash
    }

    /// Signs the specified 'message` with this private key.
    pub fn sign(&self, message: &[u8]) -> Signature {
        let mut ones = Vec::new();
        let mut zeros = Vec::new();

        let mut n = 0;
        let elements = message_to_elements(message);
        for element_bits in elements.iter().map(|e| e.as_int()) {
            // make sure the most significant bit is 0
            assert_eq!(element_bits & (1 << 127), 0);
            for i in 0..127 {
                if (element_bits >> i) & 1 == 1 {
                    ones.push(self.sec_keys[n]);
                } else {
                    zeros.push(self.pub_keys[n]);
                }
                n += 1;
            }
        }

        Signature { ones, zeros }
    }
}

// PUBLIC KEY IMPLEMENTATION
// ================================================================================================

impl PublicKey {
    /// Returns true if the specified signature was generated by signing the specified message
    /// with a private key corresponding to this public key.
    pub fn verify(&self, message: &[u8], sig: &Signature) -> bool {
        let mut n_zeros = 0;
        let mut n_ones = 0;
        let mut pub_keys = Vec::with_capacity(MESSAGE_BITS);
        let elements = message_to_elements(message);
        for element_bits in elements.iter().map(|e| e.as_int()) {
            // make sure the least significant bit is 0
            assert_eq!(element_bits & (1 << 127), 0);
            for i in 0..127 {
                if (element_bits >> i) & 1 == 1 {
                    if n_ones == sig.ones.len() {
                        return false;
                    }
                    pub_keys.push(Rescue128::digest(&sig.ones[n_ones]).to_elements());
                    n_ones += 1;
                } else {
                    if n_zeros == sig.zeros.len() {
                        return false;
                    }
                    pub_keys.push(sig.zeros[n_zeros]);
                    n_zeros += 1;
                }
            }
        }

        let pub_key_hash = hash_pub_keys(&pub_keys);
        *self == pub_key_hash
    }

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

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

impl Default for PublicKey {
    fn default() -> Self {
        PublicKey([BaseElement::ZERO; 2])
    }
}

impl Ord for PublicKey {
    fn cmp(&self, other: &Self) -> Ordering {
        if self.0[0] == other.0[0] {
            self.0[1].as_int().cmp(&other.0[1].as_int())
        } else {
            self.0[0].as_int().cmp(&other.0[0].as_int())
        }
    }
}

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

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

pub fn message_to_elements(message: &[u8]) -> [BaseElement; 2] {
    // reduce the message to a 32-byte value
    let hash = *blake3::hash(message).as_bytes();

    // interpret 32 bytes as two 128-bit integers
    let mut m0 = u128::from_le_bytes(hash[..16].try_into().unwrap());
    let mut m1 = u128::from_le_bytes(hash[16..].try_into().unwrap());

    // clear the most significant bit of the first value to ensure that it fits into 127 bits
    m0 = (m0 << 1) >> 1;

    // do the same thing with the second value, but also clear 8 more bits to make room for
    // checksum bits
    m1 = (m1 << 9) >> 9;

    // compute the checksum and put it into the most significant bits of the second values;
    // specifically: bit 127 is zeroed out, and 8 bits of checksum should go into bits
    // 119..127 thus, we just shift the checksum left by 119 bits and OR it with m1 (which
    // has top 9 bits zeroed out)
    let checksum = m0.count_zeros() + m1.count_zeros();
    let m1 = m1 | ((checksum as u128) << 119);

    [BaseElement::from(m0), BaseElement::from(m1)]
}

/// Reduces a list of public key elements to a single 32-byte value. The reduction is done
/// by breaking the list into two equal parts, and then updating hash state by taking turns
/// drawing elements from each list. For example, the final hash would be equivalent to:
/// hash(key[0] | key[127] | key[1] | key[128] | key[2] | key[129] ... )
/// This hashing methodology is implemented to simplify AIR design.
fn hash_pub_keys(keys: &[KeyData]) -> PublicKey {
    let mut pub_key_hash = Rescue128::new();
    pub_key_hash.update(&[BaseElement::ZERO; 4]);
    for i in 0..(MESSAGE_BITS / 2) {
        pub_key_hash.update(&keys[i]);
        pub_key_hash.update(&keys[i + MESSAGE_BITS / 2]);
    }

    PublicKey(pub_key_hash.finalize().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.

use super::{
    rescue, CYCLE_LENGTH as HASH_CYCLE_LEN, SIG_CYCLE_LENGTH as SIG_CYCLE_LEN, TRACE_WIDTH,
};
use crate::utils::{are_equal, is_binary, is_zero, not, EvaluationResult};
use core_utils::flatten_slice_elements;
use winterfell::{
    math::{fields::f128::BaseElement, FieldElement, ToElements},
    Air, AirContext, Assertion, EvaluationFrame, ProofOptions, TraceInfo,
    TransitionConstraintDegree,
};

// CONSTANTS
// ================================================================================================
const TWO: BaseElement = BaseElement::new(2);

// AGGREGATE LAMPORT PLUS SIGNATURE AIR
// ================================================================================================

#[derive(Clone)]
pub struct PublicInputs {
    pub pub_keys: Vec<[BaseElement; 2]>,
    pub messages: Vec<[BaseElement; 2]>,
}

impl ToElements<BaseElement> for PublicInputs {
    fn to_elements(&self) -> Vec<BaseElement> {
        let mut result = Vec::new();
        result.extend_from_slice(flatten_slice_elements(&self.pub_keys));
        result.extend_from_slice(flatten_slice_elements(&self.messages));
        result
    }
}

pub struct LamportAggregateAir {
    context: AirContext<BaseElement>,
    pub_keys: Vec<[BaseElement; 2]>,
    messages: Vec<[BaseElement; 2]>,
}

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

    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    fn new(trace_info: TraceInfo, pub_inputs: PublicInputs, options: ProofOptions) -> Self {
        // define degrees for all transition constraints
        let degrees = vec![
            TransitionConstraintDegree::with_cycles(2, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]), // m0 bit is binary
            TransitionConstraintDegree::with_cycles(2, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]), // m1 bit is binary
            TransitionConstraintDegree::with_cycles(
                1,
                vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN, SIG_CYCLE_LEN],
            ), // m0 accumulation
            TransitionConstraintDegree::with_cycles(
                1,
                vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN, SIG_CYCLE_LEN],
            ), // m1 accumulation
            // secret key 1 hashing
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            // secret key 2 hashing
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            // public key hashing
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]),
        ];
        assert_eq!(TRACE_WIDTH, trace_info.width());
        LamportAggregateAir {
            context: AirContext::new(trace_info, degrees, 22, options),
            pub_keys: pub_inputs.pub_keys,
            messages: pub_inputs.messages,
        }
    }

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

    fn evaluate_transition<E: FieldElement + From<Self::BaseField>>(
        &self,
        frame: &EvaluationFrame<E>,
        periodic_values: &[E],
        result: &mut [E],
    ) {
        let current = frame.current();
        let next = frame.next();
        // expected state width is 4 field elements
        debug_assert_eq!(TRACE_WIDTH, current.len());
        debug_assert_eq!(TRACE_WIDTH, next.len());

        // spit periodic values into flags and Rescue round constants
        let sig_cycle_end_flag = periodic_values[0];
        let power_of_two = periodic_values[1];
        let hash_flag = periodic_values[2];
        let ark = &periodic_values[3..];

        // evaluate the constraints
        evaluate_constraints(
            result,
            current,
            next,
            ark,
            hash_flag,
            sig_cycle_end_flag,
            power_of_two,
        );
    }

    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
        let last_cycle_step = SIG_CYCLE_LEN - 1;
        let messages = transpose(&self.messages);
        let pub_keys = transpose(&self.pub_keys);
        vec![
            // --- set assertions against the first step of every cycle: 0, 1024, 2048 etc. -------
            // message aggregators should be set to zeros
            Assertion::periodic(2, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(3, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            // for private key hasher, last 4 state register should be set to zeros
            Assertion::periodic(6, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(7, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(8, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(9, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(12, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(13, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(14, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(15, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            // for public key hasher, all registers should be set to zeros
            Assertion::periodic(16, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(17, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(18, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(19, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(20, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(21, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            // --- set assertions against the last step of every cycle: 1023, 2047, 3071 etc. -----
            // last bits of message bit registers should be set to zeros; this is because we truncate
            // message elements to 127 bits each - so, 128th bit must always be zero
            Assertion::periodic(0, last_cycle_step, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(1, last_cycle_step, SIG_CYCLE_LEN, BaseElement::ZERO),
            // message accumulator registers should be set to message element values
            Assertion::sequence(2, last_cycle_step, SIG_CYCLE_LEN, messages.0),
            Assertion::sequence(3, last_cycle_step, SIG_CYCLE_LEN, messages.1),
            // public key hasher should terminate with public key elements
            Assertion::sequence(16, last_cycle_step, SIG_CYCLE_LEN, pub_keys.0),
            Assertion::sequence(17, last_cycle_step, SIG_CYCLE_LEN, pub_keys.1),
        ]
    }

    fn get_periodic_column_values(&self) -> Vec<Vec<Self::BaseField>> {
        let mut result = vec![];

        // signature cycle mask: 1023 zeros followed by 1 one
        let mut sig_cycle_mask = vec![BaseElement::ZERO; SIG_CYCLE_LEN];
        sig_cycle_mask[SIG_CYCLE_LEN - 1] = BaseElement::ONE;
        result.push(sig_cycle_mask);

        // build powers of two column
        let mut powers_of_two = vec![BaseElement::ZERO; SIG_CYCLE_LEN];
        let mut current_power_of_two = BaseElement::ONE;
        powers_of_two[0] = BaseElement::ONE;
        for (i, value) in powers_of_two.iter_mut().enumerate().skip(1) {
            // we switch to a new power of two once every 8 steps this. is so that a
            // new power of two is available for every hash cycle
            if i % HASH_CYCLE_LEN == 0 {
                current_power_of_two *= TWO;
            }
            *value = current_power_of_two;
        }
        result.push(powers_of_two);

        // add hash cycle mask (seven ones followed by a zero), and rescue round constants
        result.push(HASH_CYCLE_MASK.to_vec());
        result.append(&mut rescue::get_round_constants());

        result
    }
}

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

#[rustfmt::skip]
fn evaluate_constraints<E: FieldElement + From<BaseElement>>(
    result: &mut [E],
    current: &[E],
    next: &[E],
    ark: &[E],
    hash_flag: E,
    sig_cycle_end_flag: E,
    power_of_two: E,
) {
    // when hash_flag = 1 (which happens on all steps except steps which are one less than a
    // multiple of 8 - e.g. all steps except for 7, 15, 23 etc.), and we are not on the last step
    // of a signature cycle make sure the contents of the first 4 registers are copied over, and
    // for other registers, Rescue constraints are applied separately for hashing secret and
    // public keys
    let flag = not(sig_cycle_end_flag) * hash_flag;
    result.agg_constraint(0, flag, are_equal(current[0], next[0]));
    result.agg_constraint(1, flag, are_equal(current[1], next[1]));
    result.agg_constraint(2, flag, are_equal(current[2], next[2]));
    result.agg_constraint(3, flag, are_equal(current[3], next[3]));
    rescue::enforce_round(&mut result[4..10],  &current[4..10],  &next[4..10],  ark, flag);
    rescue::enforce_round(&mut result[10..16], &current[10..16], &next[10..16], ark, flag);
    rescue::enforce_round(&mut result[16..22], &current[16..22], &next[16..22], ark, flag);

    // when hash_flag = 0 (which happens on steps which are one less than a multiple of 8 - e.g. 7,
    // 15, 23 etc.), and we are not on the last step of a signature cycle:
    let flag = not(sig_cycle_end_flag) * not(hash_flag);
    // make sure values inserted into registers 0 and 1 are binary
    result.agg_constraint(0, flag, is_binary(current[0]));
    result.agg_constraint(1, flag, is_binary(current[1]));
    // make sure message values were aggregated correctly in registers 2 and 3
    let next_m0 = current[2] + current[0] * power_of_two;
    result.agg_constraint(2, flag, are_equal(next_m0, next[2]));
    let next_m1 = current[3] + current[1] * power_of_two;
    result.agg_constraint(3, flag, are_equal(next_m1, next[3]));

    // registers 6..10 and 12..16 were set to zeros
    result.agg_constraint(4, flag, is_zero(next[6]));
    result.agg_constraint(5, flag, is_zero(next[7]));
    result.agg_constraint(6, flag, is_zero(next[8]));
    result.agg_constraint(7, flag, is_zero(next[9]));
    result.agg_constraint(8, flag, is_zero(next[12]));
    result.agg_constraint(9, flag, is_zero(next[13]));
    result.agg_constraint(10, flag, is_zero(next[14]));
    result.agg_constraint(11, flag, is_zero(next[15]));

    // contents of registers 20 and 21 (capacity section of public key hasher state) were
    // copied over to the next step
    result.agg_constraint(12, flag, are_equal(current[20], next[20]));
    result.agg_constraint(13, flag, are_equal(current[21], next[21]));

    // when current bit of m0 = 1, hash of private key 1 (which should be equal to public key)
    // should be injected into the hasher state for public key aggregator
    let m0_bit = current[0];
    result.agg_constraint(14, flag * m0_bit,are_equal(current[16] + current[4], next[16]));
    result.agg_constraint(15, flag * m0_bit, are_equal(current[17] + current[5], next[17]));

    // when current bit of m1 = 1, hash of private key 2 (which should be equal to public key)
    // should be injected into the hasher state for public key aggregator
    let m1_bit = current[1];
    result.agg_constraint(16, flag * m1_bit, are_equal(current[18] + current[10], next[18]));
    result.agg_constraint(17, flag * m1_bit, are_equal(current[19] + current[11], next[19]));
}

fn transpose(values: &[[BaseElement; 2]]) -> (Vec<BaseElement>, Vec<BaseElement>) {
    let n = values[0].len();
    let mut r1 = Vec::with_capacity(n);
    let mut r2 = Vec::with_capacity(n);
    for element in values {
        r1.push(element[0]);
        r2.push(element[1]);
    }
    (r1, r2)
}

// MASKS
// ================================================================================================
const HASH_CYCLE_MASK: [BaseElement; HASH_CYCLE_LEN] = [
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ONE,
    BaseElement::ZERO,
];


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

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

mod prover;
use prover::LamportAggregateProver;

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

const TRACE_WIDTH: usize = 22;
const SIG_CYCLE_LENGTH: usize = 128 * CYCLE_LENGTH; // 1024 steps

// LAMPORT MULTI-MESSAGE, MULTI-KEY, SIGNATURE EXAMPLE
// ================================================================================================
pub fn get_example(
    options: &ExampleOptions,
    num_signatures: usize,
) -> Result<Box<dyn Example>, String> {
    let (options, hash_fn) = options.to_proof_options(28, 8);

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

pub struct LamportAggregateExample<H: ElementHasher> {
    options: ProofOptions,
    pub_keys: Vec<[BaseElement; 2]>,
    messages: Vec<[BaseElement; 2]>,
    signatures: Vec<Signature>,
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> LamportAggregateExample<H> {
    pub fn new(num_signatures: usize, options: ProofOptions) -> Self {
        assert!(
            num_signatures.is_power_of_two(),
            "number of signatures must be a power of 2"
        );
        // generate private/public key pairs for the specified number of signatures
        let mut private_keys = Vec::with_capacity(num_signatures);
        let mut public_keys = Vec::with_capacity(num_signatures);
        let now = Instant::now();
        for i in 0..num_signatures {
            private_keys.push(PrivateKey::from_seed([i as u8; 32]));
            public_keys.push(private_keys[i].pub_key().to_elements());
        }
        debug!(
            "Generated {} private-public key pairs in {} ms",
            num_signatures,
            now.elapsed().as_millis()
        );

        // sign messages
        let now = Instant::now();
        let mut signatures = Vec::new();
        let mut messages = Vec::new();
        for (i, private_key) in private_keys.iter().enumerate() {
            let msg = format!("test message {i}");
            signatures.push(private_key.sign(msg.as_bytes()));
            messages.push(message_to_elements(msg.as_bytes()));
        }
        debug!(
            "Signed {} messages in {} ms",
            num_signatures,
            now.elapsed().as_millis()
        );

        // verify signature
        let now = Instant::now();
        let mut pub_keys = Vec::new();
        for (i, signature) in signatures.iter().enumerate() {
            let pk = private_keys[i].pub_key();
            pub_keys.push(pk.to_elements());
            let msg = format!("test message {i}");
            assert!(pk.verify(msg.as_bytes(), signature));
        }
        debug!(
            "Verified {} signature in {} ms",
            num_signatures,
            now.elapsed().as_millis()
        );

        LamportAggregateExample {
            options,
            pub_keys,
            messages,
            signatures,
            _hasher: PhantomData,
        }
    }
}

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

impl<H: ElementHasher> Example for LamportAggregateExample<H>
where
    H: ElementHasher<BaseField = BaseElement>,
{
    fn prove(&self) -> StarkProof {
        // generate the execution trace
        debug!(
            "Generating proof for verifying {} Lamport+ signatures \n\
            ---------------------",
            self.signatures.len(),
        );

        // create a prover
        let prover =
            LamportAggregateProver::<H>::new(&self.pub_keys, &self.messages, self.options.clone());

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

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

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

    fn verify_with_wrong_inputs(&self, proof: StarkProof) -> Result<(), VerifierError> {
        let mut pub_keys = self.pub_keys.clone();
        pub_keys.swap(0, 1);
        let pub_inputs = PublicInputs {
            pub_keys,
            messages: self.messages.clone(),
        };
        winterfell::verify::<LamportAggregateAir, H, DefaultRandomCoin<H>>(proof, pub_inputs)
    }
}