// 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::{
    get_power_series, rescue, BaseElement, DefaultRandomCoin, ElementHasher, FieldElement,
    LamportAggregateAir, PhantomData, ProofOptions, Prover, PublicInputs, Signature, StarkField,
    TraceTable, CYCLE_LENGTH, NUM_HASH_ROUNDS, SIG_CYCLE_LENGTH, TRACE_WIDTH,
};

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

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

const TWO: BaseElement = BaseElement::new(2);
const ZERO_KEY: [BaseElement; 2] = [BaseElement::ZERO, BaseElement::ZERO];

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

struct SignatureInfo {
    m0: u128,
    m1: u128,
    key_schedule: KeySchedule,
}

struct KeySchedule {
    sec_keys1: Vec<[BaseElement; 2]>,
    sec_keys2: Vec<[BaseElement; 2]>,
    pub_keys1: Vec<[BaseElement; 2]>,
    pub_keys2: Vec<[BaseElement; 2]>,
}

// LAMPORT PROVER
// ================================================================================================

pub struct LamportAggregateProver<H: ElementHasher> {
    pub_inputs: PublicInputs,
    options: ProofOptions,
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> LamportAggregateProver<H> {
    pub fn new(
        pub_keys: &[[BaseElement; 2]],
        messages: &[[BaseElement; 2]],
        options: ProofOptions,
    ) -> Self {
        let pub_inputs = PublicInputs {
            pub_keys: pub_keys.to_vec(),
            messages: messages.to_vec(),
        };
        Self {
            pub_inputs,
            options,
            _hasher: PhantomData,
        }
    }

    pub fn build_trace(
        &self,
        messages: &[[BaseElement; 2]],
        signatures: &[Signature],
    ) -> TraceTable<BaseElement> {
        // allocate memory to hold the trace table
        let trace_length = SIG_CYCLE_LENGTH * messages.len();
        let mut trace = TraceTable::new(TRACE_WIDTH, trace_length);

        let powers_of_two = get_power_series(TWO, 128);

        trace.fragments(SIG_CYCLE_LENGTH).for_each(|mut sig_trace| {
            let i = sig_trace.index();
            let sig_info = build_sig_info(&messages[i], &signatures[i]);
            sig_trace.fill(
                |state| {
                    init_sig_verification_state(&sig_info, state);
                },
                |step, state| {
                    update_sig_verification_state(step, &sig_info, &powers_of_two, state);
                },
            );
        });

        trace
    }
}

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

    fn get_pub_inputs(&self, _trace: &Self::Trace) -> PublicInputs {
        self.pub_inputs.clone()
    }

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

// TRACE INITIALIZATION
// ================================================================================================

fn init_sig_verification_state(sig_info: &SignatureInfo, state: &mut [BaseElement]) {
    // message accumulators
    state[0] = BaseElement::new(sig_info.m0 & 1); // m0 bits
    state[1] = BaseElement::new(sig_info.m1 & 1); // m1 bits
    state[2] = BaseElement::ZERO; // m0 accumulator
    state[3] = BaseElement::ZERO; // m1 accumulator

    // secret key 1 hashing
    state[4] = sig_info.key_schedule.sec_keys1[0][0];
    state[5] = sig_info.key_schedule.sec_keys1[0][1];
    state[6] = BaseElement::ZERO;
    state[7] = BaseElement::ZERO;
    state[8] = BaseElement::ZERO; // capacity
    state[9] = BaseElement::ZERO; // capacity

    // secret key 2 hashing
    state[10] = sig_info.key_schedule.sec_keys2[0][0];
    state[11] = sig_info.key_schedule.sec_keys2[0][1];
    state[12] = BaseElement::ZERO;
    state[13] = BaseElement::ZERO;
    state[14] = BaseElement::ZERO; // capacity
    state[15] = BaseElement::ZERO; // capacity

    // public key hashing
    state[16] = BaseElement::ZERO;
    state[17] = BaseElement::ZERO;
    state[18] = BaseElement::ZERO;
    state[19] = BaseElement::ZERO;
    state[20] = BaseElement::ZERO; // capacity
    state[21] = BaseElement::ZERO; // capacity
}

// TRANSITION FUNCTION
// ================================================================================================

fn update_sig_verification_state(
    step: usize,
    sig_info: &SignatureInfo,
    powers_of_two: &[BaseElement],
    state: &mut [BaseElement],
) {
    // determine which cycle we are in and also where in the cycle we are
    let cycle_num = step / CYCLE_LENGTH;
    let cycle_step = step % CYCLE_LENGTH;

    // break the state into logical parts
    let (msg_acc_state, rest) = state.split_at_mut(4);
    let (sec_key_1_hash, rest) = rest.split_at_mut(6);
    let (sec_key_2_hash, pub_key_hash) = rest.split_at_mut(6);

    if cycle_step < NUM_HASH_ROUNDS {
        // for the first 7 steps in each cycle apply Rescue round function to
        // registers where keys are hashed; all other registers retain their values
        rescue::apply_round(sec_key_1_hash, cycle_step);
        rescue::apply_round(sec_key_2_hash, cycle_step);
        rescue::apply_round(pub_key_hash, cycle_step);
    } else {
        let m0_bit = msg_acc_state[0];
        let m1_bit = msg_acc_state[1];

        // copy next set of public keys into the registers computing hash of the public key
        update_pub_key_hash(
            pub_key_hash,
            m0_bit,
            m1_bit,
            sec_key_1_hash,
            sec_key_2_hash,
            &sig_info.key_schedule.pub_keys1[cycle_num],
            &sig_info.key_schedule.pub_keys2[cycle_num],
        );

        // copy next set of private keys into the registers computing private key hashes
        init_hash_state(
            sec_key_1_hash,
            &sig_info.key_schedule.sec_keys1[cycle_num + 1],
        );
        init_hash_state(
            sec_key_2_hash,
            &sig_info.key_schedule.sec_keys2[cycle_num + 1],
        );

        // update message accumulator with the next set of message bits
        apply_message_acc(
            msg_acc_state,
            sig_info.m0,
            sig_info.m1,
            cycle_num,
            powers_of_two[cycle_num],
        );
    }
}

fn apply_message_acc(
    state: &mut [BaseElement],
    m0: u128,
    m1: u128,
    cycle_num: usize,
    power_of_two: BaseElement,
) {
    let m0_bit = state[0];
    let m1_bit = state[1];

    state[0] = BaseElement::from((m0 >> (cycle_num + 1)) & 1);
    state[1] = BaseElement::from((m1 >> (cycle_num + 1)) & 1);
    state[2] += power_of_two * m0_bit;
    state[3] += power_of_two * m1_bit;
}

fn init_hash_state(state: &mut [BaseElement], values: &[BaseElement; 2]) {
    state[0] = values[0];
    state[1] = values[1];
    state[2] = BaseElement::ZERO;
    state[3] = BaseElement::ZERO;
    state[4] = BaseElement::ZERO;
    state[5] = BaseElement::ZERO;
}

fn update_pub_key_hash(
    state: &mut [BaseElement],
    m0_bit: BaseElement,
    m1_bit: BaseElement,
    sec_key1_hash: &[BaseElement],
    sec_key2_hash: &[BaseElement],
    pub_key1: &[BaseElement],
    pub_key2: &[BaseElement],
) {
    if m0_bit == FieldElement::ONE {
        state[0] += sec_key1_hash[0];
        state[1] += sec_key1_hash[1];
    } else {
        state[0] += pub_key1[0];
        state[1] += pub_key1[1];
    }

    if m1_bit == FieldElement::ONE {
        state[2] += sec_key2_hash[0];
        state[3] += sec_key2_hash[1];
    } else {
        state[2] += pub_key2[0];
        state[3] += pub_key2[1];
    }
}

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

fn build_sig_info(msg: &[BaseElement; 2], sig: &Signature) -> SignatureInfo {
    let m0 = msg[0].as_int();
    let m1 = msg[1].as_int();
    let key_schedule = build_key_schedule(m0, m1, sig);
    SignatureInfo {
        m0,
        m1,
        key_schedule,
    }
}

/// Transforms signature into 4 vectors of keys such that keys 0..127 and 127..254 end up in
/// different vectors; keys that are missing from the signature are replaced with a zeros.
fn build_key_schedule(m0: u128, m1: u128, sig: &Signature) -> KeySchedule {
    let mut n_ones = 0;
    let mut n_zeros = 0;
    let mut result = KeySchedule {
        sec_keys1: vec![ZERO_KEY; 128],
        sec_keys2: vec![ZERO_KEY; 128],
        pub_keys1: vec![ZERO_KEY; 128],
        pub_keys2: vec![ZERO_KEY; 128],
    };

    for i in 0..127 {
        if (m0 >> i) & 1 == 1 {
            result.sec_keys1[i] = sig.ones[n_ones];
            n_ones += 1;
        } else {
            result.pub_keys1[i] = sig.zeros[n_zeros];
            n_zeros += 1;
        }
    }

    for i in 0..127 {
        if (m1 >> i) & 1 == 1 {
            result.sec_keys2[i] = sig.ones[n_ones];
            n_ones += 1;
        } else {
            result.pub_keys2[i] = sig.zeros[n_zeros];
            n_zeros += 1;
        }
    }

    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::{
    super::rescue, HASH_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 winterfell::{
    math::{fields::f128::BaseElement, FieldElement, StarkField, ToElements},
    Air, AirContext, Assertion, EvaluationFrame, ProofOptions, TraceInfo,
    TransitionConstraintDegree,
};

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

// THRESHOLD LAMPORT PLUS SIGNATURE AIR
// ================================================================================================

#[derive(Clone)]
pub struct PublicInputs {
    pub pub_key_root: [BaseElement; 2],
    pub num_pub_keys: usize,
    pub num_signatures: usize,
    pub message: [BaseElement; 2],
}

impl ToElements<BaseElement> for PublicInputs {
    fn to_elements(&self) -> Vec<BaseElement> {
        let mut result = self.pub_key_root.to_vec();
        result.push(BaseElement::from(self.num_pub_keys as u64));
        result.push(BaseElement::from(self.num_signatures as u64));
        result.extend_from_slice(&self.message);
        result
    }
}

pub struct LamportThresholdAir {
    context: AirContext<BaseElement>,
    pub_key_root: [BaseElement; 2],
    num_pub_keys: usize,
    num_signatures: usize,
    message: [BaseElement; 2],
}

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

    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    #[rustfmt::skip]
    fn new(trace_info: TraceInfo, pub_inputs: PublicInputs, options: ProofOptions) -> Self {
        // define degrees for all transition constraints
        let degrees = vec![
            // 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]),
            // merkle path verification
            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]),
            // merkle path index
            TransitionConstraintDegree::with_cycles(2, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN]), // index bit is binary
            TransitionConstraintDegree::with_cycles(1, vec![HASH_CYCLE_LEN, SIG_CYCLE_LEN, SIG_CYCLE_LEN]), // index accumulator
            // signature count
            TransitionConstraintDegree::with_cycles(2, vec![SIG_CYCLE_LEN]), // sig flag is binary
            TransitionConstraintDegree::with_cycles(1, vec![SIG_CYCLE_LEN]), // sig counter
            TransitionConstraintDegree::with_cycles(2, vec![SIG_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(2, vec![SIG_CYCLE_LEN]),
        ];
        assert_eq!(TRACE_WIDTH, trace_info.width());
        LamportThresholdAir {
            context: AirContext::new(trace_info, degrees, 26, options),
            pub_key_root: pub_inputs.pub_key_root,
            num_pub_keys: pub_inputs.num_pub_keys,
            num_signatures: pub_inputs.num_signatures,
            message: pub_inputs.message,
        }
    }

    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 m0_bit = periodic_values[2];
        let m1_bit = periodic_values[3];
        let hash_flag = periodic_values[4];
        let ark = &periodic_values[5..];

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

    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
        // ----- assertions against the first step of every cycle: 0, 1024, 2048 etc. -------------
        let mut assertions = vec![
            // for private key hasher, last 4 state register should be set to zeros
            Assertion::periodic(2, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(3, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(4, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(5, 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(10, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(11, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            // for public key hasher, all registers should be set to zeros
            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),
            Assertion::periodic(16, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(17, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            // for merkle path verification, last 4 registers should be set to zeros
            Assertion::periodic(20, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(21, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(22, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(23, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
            // merkle path index accumulator should be initialized to zero
            Assertion::periodic(25, 0, SIG_CYCLE_LEN, BaseElement::ZERO),
        ];

        // ----- assertions against the step in every cycle when the Merkle path computation ends -
        // these steps depend on the depth of the public key Merkle tree; for example, if the Merkle
        // tree has 4 elements, then the steps are: 24, 1048, 2072, 3096
        let num_cycles = self.num_pub_keys.next_power_of_two();
        let merkle_root_offset = (num_cycles.ilog2() + 1) as usize * HASH_CYCLE_LEN;

        // distinct key indexes should be used; the sequence starts at the last index of the tree
        // (to pad the first cycle) and then wraps around and proceeds with index 0, 1, 2 etc.
        let index_list = get_index_list(num_cycles);

        assertions.extend_from_slice(&[
            Assertion::sequence(25, merkle_root_offset, SIG_CYCLE_LEN, index_list),
            // merkle path verifications should terminate with the root public key
            Assertion::periodic(18, merkle_root_offset, SIG_CYCLE_LEN, self.pub_key_root[0]),
            Assertion::periodic(19, merkle_root_offset, SIG_CYCLE_LEN, self.pub_key_root[1]),
        ]);

        // ----- assertions for the entire execution trace -----------------------------------------

        let last_step = self.trace_length() - 1;
        assertions.extend_from_slice(&[
            // signature counter starts at zero and terminates with the expected count of signatures
            Assertion::single(27, 0, BaseElement::ZERO),
            Assertion::single(27, last_step, BaseElement::from(self.num_signatures as u64)),
            // the first public key for merkle path verification should be a zero key (it is only
            // used for padding)
            Assertion::single(18, 0, BaseElement::ZERO),
            Assertion::single(19, 0, BaseElement::ZERO),
        ]);
        assertions
    }

    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);

        // build message bit columns m0 and m1
        let m0 = self.message[0].as_int();
        let m1 = self.message[1].as_int();
        let mut m0_bits = Vec::with_capacity(SIG_CYCLE_LEN);
        let mut m1_bits = Vec::with_capacity(SIG_CYCLE_LEN);
        for i in 0..SIG_CYCLE_LEN {
            let cycle_num = i / HASH_CYCLE_LEN;
            m0_bits.push(BaseElement::from((m0 >> cycle_num) & 1));
            m1_bits.push(BaseElement::from((m1 >> cycle_num) & 1));
        }
        result.push(m0_bits);
        result.push(m1_bits);

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

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

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

#[rustfmt::skip]
#[allow(clippy::too_many_arguments)]
fn evaluate_constraints<E: FieldElement + From<BaseElement>>(
    result: &mut [E],
    current: &[E],
    next: &[E],
    ark: &[E],
    hash_flag: E,
    sig_cycle_end_flag: E,
    m0_bit: E,
    m1_bit: 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 registers 24 and 25 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;
    rescue::enforce_round(&mut result[..6],  &current[..6],  &next[..6],  ark, flag);
    rescue::enforce_round(&mut result[6..12], &current[6..12], &next[6..12], ark, flag);
    rescue::enforce_round(&mut result[12..18], &current[12..18], &next[12..18], ark, flag);
    rescue::enforce_round(&mut result[18..24], &current[18..24], &next[18..24], ark, flag);
    result.agg_constraint(24, flag, are_equal(current[24], next[24]));
    result.agg_constraint(25, flag, are_equal(current[25], next[25]));

    // 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);
    // registers 2..6 and 8..12 were set to zeros
    result.agg_constraint(0, flag, is_zero(next[2]));
    result.agg_constraint(1, flag, is_zero(next[3]));
    result.agg_constraint(2, flag, is_zero(next[4]));
    result.agg_constraint(3, flag, is_zero(next[5]));
    result.agg_constraint(4, flag, is_zero(next[8]));
    result.agg_constraint(5, flag, is_zero(next[9]));
    result.agg_constraint(6, flag, is_zero(next[10]));
    result.agg_constraint(7, flag, is_zero(next[11]));

    // contents of registers 16 and 17 (capacity section of public key hasher state) were
    // copied over to the next step
    result.agg_constraint(8, flag, are_equal(current[16], next[16]));
    result.agg_constraint(9, flag, are_equal(current[17], next[17]));

    // 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
    result.agg_constraint(10, flag * m0_bit,are_equal(current[12] + current[0], next[12]));
    result.agg_constraint(11, flag * m0_bit, are_equal(current[13] + current[1], next[13]));

    // 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
    result.agg_constraint(16, flag * m1_bit, are_equal(current[14] + current[6], next[14]));
    result.agg_constraint(17, flag * m1_bit, are_equal(current[15] + current[7], next[15]));

    // when merkle path bit = 1, next values for registers 18 and 19 should come from
    // registers 20 and 21; but when the bit = 0, values should be copied over from
    // registers 18 and 19; registers 22 and 23 should be reset to zeros.
    let mp_bit = current[24];
    result.agg_constraint(18, flag * not(mp_bit), are_equal(current[18], next[18]));
    result.agg_constraint(19, flag * not(mp_bit), are_equal(current[19], next[19]));
    result.agg_constraint(20, flag * mp_bit, are_equal(current[18], next[20]));
    result.agg_constraint(21, flag * mp_bit, are_equal(current[19], next[21]));
    result.agg_constraint(22, flag, is_zero(next[22]));
    result.agg_constraint(23, flag, is_zero(next[23]));

    // make sure merkle path index bit is binary
    result.agg_constraint(24, flag, is_binary(current[24]));
    // make sure merkle path index aggregator is incremented correctly
    let next_index_agg = current[25] + current[24] * power_of_two;
    result.agg_constraint(25, flag, are_equal(next_index_agg, next[25]));

    // sig flag should be binary and shouldn't change during the signature cycle
    let sig_flag = current[26];
    result.agg_constraint(26, not(sig_cycle_end_flag), are_equal(sig_flag, next[26]));
    result.agg_constraint(26, sig_cycle_end_flag, is_binary(sig_flag));

    // on all steps but the last step of the signature cycle, sig count should be copied
    // over to the next step; on the last step of the signature cycle the next value of
    // sig count should be set to the previous value, plus the current value of sig flag
    result.agg_constraint(27, not(sig_cycle_end_flag), are_equal(current[27], next[27]));
    result.agg_constraint(27, sig_cycle_end_flag, are_equal(current[27] + sig_flag, next[27]));

    // when sig_count=1, public key computed during signature verification (registers 12 and 13)
    // should be copied to the beginning of Merkle path computation for the aggregated public key
    // (registers 18 and 19); this constraint should be enforced only on the last step of signature
    // verification cycle
    result.agg_constraint(28, sig_cycle_end_flag * sig_flag, are_equal(current[12], next[18]));
    result.agg_constraint(29, sig_cycle_end_flag * sig_flag, are_equal(current[13], next[19]));
}

fn get_index_list(num_keys: usize) -> Vec<BaseElement> {
    let mut result = Vec::with_capacity(num_keys);
    result.push(BaseElement::from((num_keys - 1) as u64));
    for i in 0..(num_keys - 1) {
        result.push(BaseElement::from(i as u64));
    }
    result
}

// 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 as HASH_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 signature;
use signature::AggPublicKey;

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

mod prover;
use prover::LamportThresholdProver;

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

const TRACE_WIDTH: usize = 28;
const SIG_CYCLE_LENGTH: usize = 128 * HASH_CYCLE_LENGTH; // 1024 steps

// LAMPORT THRESHOLD SIGNATURE EXAMPLE
// ================================================================================================

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

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

pub struct LamportThresholdExample<H: ElementHasher> {
    options: ProofOptions,
    pub_key: AggPublicKey,
    signatures: Vec<(usize, Signature)>,
    message: [BaseElement; 2],
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> LamportThresholdExample<H> {
    pub fn new(num_signers: usize, options: &ExampleOptions) -> Self {
        assert!(
            (num_signers + 1).is_power_of_two(),
            "number of signers must be one less than a power of 2"
        );
        // generate private/public key pairs for the specified number of signatures
        let now = Instant::now();
        let private_keys = build_keys(num_signers);
        debug!(
            "Generated {} private-public key pairs in {} ms",
            num_signers,
            now.elapsed().as_millis()
        );
        let public_keys = private_keys.iter().map(|k| k.pub_key()).collect();

        // sign the message with the subset of previously generated keys
        let message = "test message";
        let selected_indexes = pick_random_indexes(num_signers);
        let mut signatures = Vec::new();
        for &key_index in selected_indexes.iter() {
            let signature = private_keys[key_index].sign(message.as_bytes());
            signatures.push((key_index, signature));
        }

        // build the aggregated public key
        let now = Instant::now();
        let pub_key = AggPublicKey::new(public_keys);
        debug!(
            "Built aggregated public key in {} ms",
            now.elapsed().as_millis()
        );

        let (options, _) = options.to_proof_options(28, 8);

        LamportThresholdExample {
            options,
            pub_key,
            signatures,
            message: message_to_elements(message.as_bytes()),
            _hasher: PhantomData,
        }
    }
}

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

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

        // create a prover
        let prover = LamportThresholdProver::<H>::new(
            &self.pub_key,
            self.message,
            &self.signatures,
            self.options.clone(),
        );

        // generate execution trace
        let now = Instant::now();
        let trace = prover.build_trace(&self.pub_key, self.message, &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_key_root: self.pub_key.root().to_elements(),
            num_pub_keys: self.pub_key.num_keys(),
            num_signatures: self.signatures.len(),
            message: self.message,
        };
        winterfell::verify::<LamportThresholdAir, H, DefaultRandomCoin<H>>(proof, pub_inputs)
    }

    fn verify_with_wrong_inputs(&self, proof: StarkProof) -> Result<(), VerifierError> {
        let pub_inputs = PublicInputs {
            pub_key_root: self.pub_key.root().to_elements(),
            num_pub_keys: self.pub_key.num_keys(),
            num_signatures: self.signatures.len() + 1,
            message: self.message,
        };
        winterfell::verify::<LamportThresholdAir, H, DefaultRandomCoin<H>>(proof, pub_inputs)
    }
}

// HELPER FUNCTIONS
// ================================================================================================
fn build_keys(num_keys: usize) -> Vec<PrivateKey> {
    let mut result = Vec::with_capacity(num_keys);
    for i in 0..num_keys {
        result.push(PrivateKey::from_seed([i as u8; 32]));
    }
    result.sort_by_key(|k| k.pub_key());
    result
}

fn pick_random_indexes(num_keys: usize) -> Vec<usize> {
    let num_selected_keys = num_keys * 2 / 3;
    let mut result = Vec::with_capacity(num_selected_keys);
    // TODO: change to actual random selection
    for i in 0..num_selected_keys {
        result.push(i);
    }
    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::{
    get_power_series, rescue, AggPublicKey, BaseElement, DefaultRandomCoin, ElementHasher,
    FieldElement, LamportThresholdAir, PhantomData, ProofOptions, Prover, PublicInputs, Signature,
    StarkField, TraceTable, HASH_CYCLE_LENGTH, NUM_HASH_ROUNDS, SIG_CYCLE_LENGTH, TRACE_WIDTH,
};
use std::collections::HashMap;

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

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

const TWO: BaseElement = BaseElement::new(2);
const ZERO_KEY: [BaseElement; 2] = [BaseElement::ZERO, BaseElement::ZERO];

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

struct SignatureInfo {
    m0: u128,
    m1: u128,
    pub_key: [BaseElement; 2],
    key_schedule: KeySchedule,
    key_index: u128,
    key_path: Vec<rescue::Hash>,
    sig_flag: BaseElement,
    sig_count: BaseElement,
}

struct KeySchedule {
    sec_keys1: Vec<[BaseElement; 2]>,
    sec_keys2: Vec<[BaseElement; 2]>,
    pub_keys1: Vec<[BaseElement; 2]>,
    pub_keys2: Vec<[BaseElement; 2]>,
}

// LAMPORT PROVER
// ================================================================================================

pub struct LamportThresholdProver<H: ElementHasher> {
    pub_inputs: PublicInputs,
    options: ProofOptions,
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> LamportThresholdProver<H> {
    pub fn new(
        pub_key: &AggPublicKey,
        message: [BaseElement; 2],
        signatures: &[(usize, Signature)],
        options: ProofOptions,
    ) -> Self {
        let pub_inputs = PublicInputs {
            pub_key_root: pub_key.root().to_elements(),
            num_pub_keys: pub_key.num_keys(),
            num_signatures: signatures.len(),
            message,
        };
        Self {
            pub_inputs,
            options,
            _hasher: PhantomData,
        }
    }

    pub fn build_trace(
        &self,
        pub_key: &AggPublicKey,
        message: [BaseElement; 2],
        signatures: &[(usize, Signature)],
    ) -> TraceTable<BaseElement> {
        // allocate memory to hold the trace table
        let num_cycles = pub_key.num_keys().next_power_of_two();
        let trace_length = SIG_CYCLE_LENGTH * num_cycles;
        let mut trace = TraceTable::new(TRACE_WIDTH, trace_length);

        let powers_of_two = get_power_series(TWO, 128);

        // transform a list of signatures into a hashmap; this way we can look up signature
        // by index of the corresponding public key
        let mut signature_map = HashMap::new();
        for (i, sig) in signatures {
            signature_map.insert(i, sig);
        }

        // build a map of signature indexes to the running sum of valid signatures
        let mut sig_count = vec![0];
        for i in 1..num_cycles {
            match signature_map.get(&(i - 1)) {
                Some(_) => sig_count.push(sig_count[i - 1] + 1),
                None => sig_count.push(sig_count[i - 1]),
            }
        }

        // create a dummy signature; this will be used in place of signatures for keys
        // which did not sign the message
        let zero_sig = Signature {
            ones: vec![[BaseElement::ZERO; 2]; 254],
            zeros: vec![[BaseElement::ZERO; 2]; 254],
        };

        // iterate over all leaves of the aggregated public key; and if a signature exists for the
        // corresponding individual public key, use it go generate signature verification trace;
        // otherwise, use zero signature;
        trace.fragments(SIG_CYCLE_LENGTH).for_each(|mut sig_trace| {
            let i = sig_trace.index();
            let sig_info = match signature_map.get(&i) {
                Some(sig) => build_sig_info(i, &message, sig, 1, pub_key, sig_count[i]),
                None => build_sig_info(i, &message, &zero_sig, 0, pub_key, sig_count[i]),
            };

            sig_trace.fill(
                |state| {
                    init_sig_verification_state(&sig_info, state);
                },
                |step, state| {
                    update_sig_verification_state(step, &sig_info, &powers_of_two, state);
                },
            );
        });

        trace
    }
}

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

    fn get_pub_inputs(&self, _trace: &Self::Trace) -> PublicInputs {
        self.pub_inputs.clone()
    }

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

// TRACE INITIALIZATION
// ================================================================================================

fn init_sig_verification_state(sig_info: &SignatureInfo, state: &mut [BaseElement]) {
    // secret key 1 hashing
    state[0] = sig_info.key_schedule.sec_keys1[0][0];
    state[1] = sig_info.key_schedule.sec_keys1[0][1];
    state[2] = BaseElement::ZERO;
    state[3] = BaseElement::ZERO;
    state[4] = BaseElement::ZERO; // capacity
    state[5] = BaseElement::ZERO; // capacity
                                  // secret key 2 hashing
    state[6] = sig_info.key_schedule.sec_keys2[0][0];
    state[7] = sig_info.key_schedule.sec_keys2[0][1];
    state[8] = BaseElement::ZERO;
    state[9] = BaseElement::ZERO;
    state[10] = BaseElement::ZERO; // capacity
    state[11] = BaseElement::ZERO; // capacity
                                   // public key hashing
    state[12] = BaseElement::ZERO;
    state[13] = BaseElement::ZERO;
    state[14] = BaseElement::ZERO;
    state[15] = BaseElement::ZERO;
    state[16] = BaseElement::ZERO; // capacity
    state[17] = BaseElement::ZERO; // capacity
                                   // merkle path verification
    state[18] = sig_info.pub_key[0];
    state[19] = sig_info.pub_key[1];
    state[20] = BaseElement::ZERO;
    state[21] = BaseElement::ZERO;
    state[22] = BaseElement::ZERO; // capacity
    state[23] = BaseElement::ZERO; // capacity
    state[24] = BaseElement::new(sig_info.key_index & 1); // index bits
    state[25] = BaseElement::ZERO; // index accumulator
                                   // signature counter
    state[26] = sig_info.sig_flag; // signature flag
    state[27] = sig_info.sig_count; // signature count
}

// TRANSITION FUNCTION
// ================================================================================================

fn update_sig_verification_state(
    step: usize,
    sig_info: &SignatureInfo,
    powers_of_two: &[BaseElement],
    state: &mut [BaseElement],
) {
    // determine which cycle we are in and also where in the cycle we are
    let cycle_num = (step % SIG_CYCLE_LENGTH) / HASH_CYCLE_LENGTH;
    let cycle_step = (step % SIG_CYCLE_LENGTH) % HASH_CYCLE_LENGTH;

    // break the state into logical parts; we don't need to do anything with sig_count part
    // because values for these registers are set in the initial state and don't change
    // during the cycle
    let (sec_key_1_hash, rest) = state.split_at_mut(6);
    let (sec_key_2_hash, rest) = rest.split_at_mut(6);
    let (pub_key_hash, rest) = rest.split_at_mut(6);
    let (merkle_path_hash, rest) = rest.split_at_mut(6);
    let (merkle_path_idx, _sig_count) = rest.split_at_mut(2);

    if cycle_step < NUM_HASH_ROUNDS {
        // for the first 7 steps in each hash cycle apply Rescue round function to
        // registers where keys are hashed; all other registers retain their values
        rescue::apply_round(sec_key_1_hash, cycle_step);
        rescue::apply_round(sec_key_2_hash, cycle_step);
        rescue::apply_round(pub_key_hash, cycle_step);
        rescue::apply_round(merkle_path_hash, cycle_step);
    } else {
        // for the 8th step of very cycle do the following:

        let m0_bit = BaseElement::from((sig_info.m0 >> cycle_num) & 1);
        let m1_bit = BaseElement::from((sig_info.m1 >> cycle_num) & 1);
        let mp_bit = merkle_path_idx[0];

        // copy next set of public keys into the registers computing hash of the public key
        update_pub_key_hash(
            pub_key_hash,
            m0_bit,
            m1_bit,
            sec_key_1_hash,
            sec_key_2_hash,
            &sig_info.key_schedule.pub_keys1[cycle_num],
            &sig_info.key_schedule.pub_keys2[cycle_num],
        );

        // copy next set of private keys into the registers computing private key hashes
        init_hash_state(
            sec_key_1_hash,
            &sig_info.key_schedule.sec_keys1[cycle_num + 1],
        );
        init_hash_state(
            sec_key_2_hash,
            &sig_info.key_schedule.sec_keys2[cycle_num + 1],
        );

        // update merkle path index accumulator with the next index bit
        update_merkle_path_index(
            merkle_path_idx,
            sig_info.key_index,
            cycle_num,
            powers_of_two[cycle_num],
        );
        // prepare Merkle path hashing registers for hashing of the next node
        update_merkle_path_hash(merkle_path_hash, mp_bit, cycle_num, &sig_info.key_path);
    }
}

fn init_hash_state(state: &mut [BaseElement], values: &[BaseElement; 2]) {
    state[0] = values[0];
    state[1] = values[1];
    state[2] = BaseElement::ZERO;
    state[3] = BaseElement::ZERO;
    state[4] = BaseElement::ZERO;
    state[5] = BaseElement::ZERO;
}

fn update_pub_key_hash(
    state: &mut [BaseElement],
    m0_bit: BaseElement,
    m1_bit: BaseElement,
    sec_key1_hash: &[BaseElement],
    sec_key2_hash: &[BaseElement],
    pub_key1: &[BaseElement],
    pub_key2: &[BaseElement],
) {
    if m0_bit == FieldElement::ONE {
        state[0] += sec_key1_hash[0];
        state[1] += sec_key1_hash[1];
    } else {
        state[0] += pub_key1[0];
        state[1] += pub_key1[1];
    }

    if m1_bit == FieldElement::ONE {
        state[2] += sec_key2_hash[0];
        state[3] += sec_key2_hash[1];
    } else {
        state[2] += pub_key2[0];
        state[3] += pub_key2[1];
    }
}

fn update_merkle_path_hash(
    state: &mut [BaseElement],
    index_bit: BaseElement,
    cycle_num: usize,
    key_path: &[rescue::Hash],
) {
    let h1 = state[0];
    let h2 = state[1];
    let cycle_num = (cycle_num + 1) % key_path.len();
    let path_node = key_path[cycle_num].to_elements();
    if index_bit == BaseElement::ONE {
        state[0] = path_node[0];
        state[1] = path_node[1];
        state[2] = h1;
        state[3] = h2;
    } else {
        state[0] = h1;
        state[1] = h2;
        state[2] = path_node[0];
        state[3] = path_node[1];
    }
    state[4] = BaseElement::ZERO;
    state[5] = BaseElement::ZERO;
}

fn update_merkle_path_index(
    state: &mut [BaseElement],
    index: u128,
    cycle_num: usize,
    power_of_two: BaseElement,
) {
    let index_bit = state[0];
    // the cycle is offset by +1 because the first node in the Merkle path is redundant and we
    // get it by hashing the public key
    state[0] = BaseElement::from((index >> (cycle_num + 1)) & 1);
    state[1] += power_of_two * index_bit;
}

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

fn build_sig_info(
    index: usize,
    msg: &[BaseElement; 2],
    sig: &Signature,
    sig_flag: u64,
    pub_key: &AggPublicKey,
    sig_count: u64,
) -> SignatureInfo {
    let m0 = msg[0].as_int();
    let m1 = msg[1].as_int();
    // we verify that the individual public key exists in the aggregated public key after
    // we've verified the signature; thus, the key index is offset by 1. That is, when
    // we verify signature for pub key 1, we verify Merkle path for pub key 0; the last
    // verification wraps around, but we don't care since the last signature is always a
    // zero signature which does not affect the count.
    let key_index = sig_index_to_key_index(index, pub_key.num_leaves());
    SignatureInfo {
        m0,
        m1,
        pub_key: pub_key.get_key(key_index).unwrap_or_default().to_elements(),
        key_schedule: build_key_schedule(m0, m1, sig),
        key_index: key_index as u128,
        key_path: pub_key.get_leaf_path(key_index),
        sig_flag: BaseElement::from(sig_flag),
        sig_count: BaseElement::from(sig_count),
    }
}

/// Transforms signature into 4 vectors of keys such that keys 0..127 and 127..254 end up in
/// different vectors; keys that are missing from the signature are replaced with a zeros.
fn build_key_schedule(m0: u128, m1: u128, sig: &Signature) -> KeySchedule {
    let mut n_ones = 0;
    let mut n_zeros = 0;
    let mut result = KeySchedule {
        sec_keys1: vec![ZERO_KEY; 128],
        sec_keys2: vec![ZERO_KEY; 128],
        pub_keys1: vec![ZERO_KEY; 128],
        pub_keys2: vec![ZERO_KEY; 128],
    };

    for i in 0..127 {
        if (m0 >> i) & 1 == 1 {
            result.sec_keys1[i] = sig.ones[n_ones];
            n_ones += 1;
        } else {
            result.pub_keys1[i] = sig.zeros[n_zeros];
            n_zeros += 1;
        }
    }

    for i in 0..127 {
        if (m1 >> i) & 1 == 1 {
            result.sec_keys2[i] = sig.ones[n_ones];
            n_ones += 1;
        } else {
            result.pub_keys2[i] = sig.zeros[n_zeros];
            n_zeros += 1;
        }
    }

    result
}

fn sig_index_to_key_index(sig_index: usize, num_cycles: usize) -> usize {
    if sig_index == 0 {
        num_cycles - 1
    } else {
        sig_index - 1
    }
}


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

use crate::{
    lamport::signature::PublicKey,
    utils::rescue::{Hash, Rescue128},
};
use winterfell::{
    crypto::MerkleTree,
    math::{fields::f128::BaseElement, FieldElement},
};

// AGGREGATED PUBLIC KEY
// ================================================================================================

pub struct AggPublicKey {
    keys: Vec<PublicKey>,
    tree: MerkleTree<Rescue128>,
}

impl AggPublicKey {
    pub fn new(mut keys: Vec<PublicKey>) -> Self {
        // sort keys in ascending order
        keys.sort();

        // convert keys to arrays of bytes; each key is hashed using Rescue hash function; the
        // initial hashing makes the AIR design a little simpler
        let mut leaves: Vec<Hash> = Vec::new();
        for key in keys.iter() {
            leaves.push(Rescue128::digest(&key.to_elements()));
        }

        // pad the list of keys with zero keys to make sure the number of leaves is greater than
        // the number of keys and is a power of two
        let num_leaves = if leaves.len().is_power_of_two() {
            (leaves.len() + 1).next_power_of_two()
        } else {
            leaves.len().next_power_of_two()
        };
        let zero_hash = Rescue128::digest(&[BaseElement::ZERO, BaseElement::ZERO]);
        for _ in leaves.len()..num_leaves {
            leaves.push(zero_hash);
        }

        // build a Merkle tree of all leaves
        let tree = MerkleTree::new(leaves).unwrap();

        AggPublicKey { keys, tree }
    }

    /// Returns a 32-byte representation of the aggregated public key.
    pub fn root(&self) -> Hash {
        *self.tree.root()
    }

    /// Returns the number of individual keys aggregated into this key.
    pub fn num_keys(&self) -> usize {
        self.keys.len()
    }

    /// Returns number of leaves in the aggregated public key; this will always be greater
    // than the number of individual keys.
    pub fn num_leaves(&self) -> usize {
        self.tree.leaves().len()
    }

    /// Returns an individual key at the specified index, if one exists.
    pub fn get_key(&self, index: usize) -> Option<PublicKey> {
        if index < self.keys.len() {
            Some(self.keys[index])
        } else {
            None
        }
    }

    /// Returns a Merkle path to the specified leaf.
    pub fn get_leaf_path(&self, index: usize) -> Vec<Hash> {
        self.tree.prove(index).unwrap()
    }
}


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

use super::{rescue, BaseElement, FieldElement, HASH_CYCLE_LEN, HASH_STATE_WIDTH, TRACE_WIDTH};
use crate::utils::{are_equal, is_binary, is_zero, not, EvaluationResult};
use winterfell::{
    math::ToElements, Air, AirContext, Assertion, EvaluationFrame, ProofOptions, TraceInfo,
    TransitionConstraintDegree,
};

// MERKLE PATH VERIFICATION AIR
// ================================================================================================

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

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

pub struct MerkleAir {
    context: AirContext<BaseElement>,
    tree_root: [BaseElement; 2],
}

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

    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    fn new(trace_info: TraceInfo, pub_inputs: PublicInputs, options: ProofOptions) -> Self {
        let degrees = vec![
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN]),
            TransitionConstraintDegree::with_cycles(5, vec![HASH_CYCLE_LEN]),
            TransitionConstraintDegree::new(2),
        ];
        assert_eq!(TRACE_WIDTH, trace_info.width());
        MerkleAir {
            context: AirContext::new(trace_info, degrees, 4, options),
            tree_root: pub_inputs.tree_root,
        }
    }

    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());

        // split periodic values into masks and Rescue round constants
        let hash_flag = periodic_values[0];
        let ark = &periodic_values[1..];

        // when hash_flag = 1, constraints for Rescue round are enforced
        rescue::enforce_round(
            result,
            &current[..HASH_STATE_WIDTH],
            &next[..HASH_STATE_WIDTH],
            ark,
            hash_flag,
        );

        // when hash_flag = 0, make sure accumulated hash is placed in the right place in the hash
        // state for the next round of hashing. Specifically: when index bit = 0 accumulated hash
        // must go into registers [0, 1], and when index bit = 0, it must go into registers [2, 3]
        let hash_init_flag = not(hash_flag);
        let bit = next[6];
        let not_bit = not(bit);
        result.agg_constraint(0, hash_init_flag, not_bit * are_equal(current[0], next[0]));
        result.agg_constraint(1, hash_init_flag, not_bit * are_equal(current[1], next[1]));
        result.agg_constraint(2, hash_init_flag, bit * are_equal(current[0], next[2]));
        result.agg_constraint(3, hash_init_flag, bit * are_equal(current[1], next[3]));

        // make sure capacity registers of the hash state are reset to zeros
        result.agg_constraint(4, hash_init_flag, is_zero(next[4]));
        result.agg_constraint(5, hash_init_flag, is_zero(next[5]));

        // finally, we always enforce that values in the bit register must be binary
        result[6] = is_binary(current[6]);
    }

    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
        // assert that Merkle path resolves to the tree root, and that hash capacity
        // registers (registers 4 and 5) are reset to ZERO every 8 steps
        let last_step = self.trace_length() - 1;
        vec![
            Assertion::single(0, last_step, self.tree_root[0]),
            Assertion::single(1, last_step, self.tree_root[1]),
            Assertion::periodic(4, 0, HASH_CYCLE_LEN, BaseElement::ZERO),
            Assertion::periodic(5, 0, HASH_CYCLE_LEN, BaseElement::ZERO),
        ]
    }

    fn get_periodic_column_values(&self) -> Vec<Vec<Self::BaseField>> {
        let mut result = vec![HASH_CYCLE_MASK.to_vec()];
        result.append(&mut rescue::get_round_constants());
        result
    }
}

// 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 crate::utils::rescue::{
    self, CYCLE_LENGTH as HASH_CYCLE_LEN, NUM_ROUNDS as NUM_HASH_ROUNDS,
    STATE_WIDTH as HASH_STATE_WIDTH,
};
use crate::{
    utils::rescue::{Hash, Rescue128},
    Blake3_192, Blake3_256, Example, ExampleOptions, HashFunction, Sha3_256,
};
use core::marker::PhantomData;
use log::debug;
use rand_utils::{rand_value, rand_vector};
use std::time::Instant;
use winterfell::{
    crypto::{DefaultRandomCoin, Digest, ElementHasher, MerkleTree},
    math::{fields::f128::BaseElement, FieldElement, StarkField},
    ProofOptions, Prover, StarkProof, Trace, TraceTable, VerifierError,
};

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

mod prover;
use prover::MerkleProver;

#[cfg(test)]
mod tests;

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

const TRACE_WIDTH: usize = 7;

// MERKLE AUTHENTICATION PATH EXAMPLE
// ================================================================================================
pub fn get_example(
    options: &ExampleOptions,
    tree_depth: 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(MerkleExample::<Blake3_192>::new(
            tree_depth, options,
        ))),
        HashFunction::Blake3_256 => Ok(Box::new(MerkleExample::<Blake3_256>::new(
            tree_depth, options,
        ))),
        HashFunction::Sha3_256 => Ok(Box::new(MerkleExample::<Sha3_256>::new(
            tree_depth, options,
        ))),
        _ => Err("The specified hash function cannot be used with this example.".to_string()),
    }
}

pub struct MerkleExample<H: ElementHasher> {
    options: ProofOptions,
    tree_root: Hash,
    value: [BaseElement; 2],
    index: usize,
    path: Vec<Hash>,
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> MerkleExample<H> {
    pub fn new(tree_depth: usize, options: ProofOptions) -> Self {
        assert!(
            (tree_depth + 1).is_power_of_two(),
            "tree depth must be one less than a power of 2"
        );
        let value = [BaseElement::new(42), BaseElement::new(43)];
        let index =
            (rand_value::<BaseElement>().as_int() % u128::pow(2, tree_depth as u32)) as usize;

        // build Merkle tree of the specified depth
        let now = Instant::now();
        let tree = build_merkle_tree(tree_depth, value, index);
        debug!(
            "Built Merkle tree of depth {} in {} ms",
            tree_depth,
            now.elapsed().as_millis(),
        );

        // compute Merkle path form the leaf specified by the index
        let now = Instant::now();
        let path = tree.prove(index).unwrap();
        debug!(
            "Computed Merkle path from leaf {} to root {} in {} ms",
            index,
            hex::encode(tree.root().as_bytes()),
            now.elapsed().as_millis(),
        );

        MerkleExample {
            options,
            tree_root: *tree.root(),
            value,
            index,
            path,
            _hasher: PhantomData,
        }
    }
}

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

impl<H: ElementHasher> Example for MerkleExample<H>
where
    H: ElementHasher<BaseField = BaseElement>,
{
    fn prove(&self) -> StarkProof {
        // generate the execution trace
        debug!(
            "Generating proof for proving membership in a Merkle tree of depth {}\n\
            ---------------------",
            self.path.len()
        );
        // create the prover
        let prover = MerkleProver::<H>::new(self.options.clone());

        // generate the execution trace
        let now = Instant::now();
        let trace = prover.build_trace(self.value, &self.path, self.index);
        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 {
            tree_root: self.tree_root.to_elements(),
        };
        winterfell::verify::<MerkleAir, H, DefaultRandomCoin<H>>(proof, pub_inputs)
    }

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

// HELPER FUNCTIONS
// ================================================================================================
fn build_merkle_tree(depth: usize, value: [BaseElement; 2], index: usize) -> MerkleTree<Rescue128> {
    let num_leaves = usize::pow(2, depth as u32);
    let leaf_elements: Vec<BaseElement> = rand_vector(num_leaves * 2);
    let mut leaves = Vec::new();
    for i in (0..leaf_elements.len()).step_by(2) {
        leaves.push(Hash::new(leaf_elements[i], leaf_elements[i + 1]));
    }

    leaves[index] = Rescue128::digest(&value);
    MerkleTree::new(leaves).unwrap()
}


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

use super::{
    rescue, BaseElement, DefaultRandomCoin, ElementHasher, FieldElement, MerkleAir, PhantomData,
    ProofOptions, Prover, PublicInputs, Trace, TraceTable, HASH_CYCLE_LEN, HASH_STATE_WIDTH,
    NUM_HASH_ROUNDS, TRACE_WIDTH,
};

// MERKLE PROVER
// ================================================================================================

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

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

    pub fn build_trace(
        &self,
        value: [BaseElement; 2],
        branch: &[rescue::Hash],
        index: usize,
    ) -> TraceTable<BaseElement> {
        // allocate memory to hold the trace table
        let trace_length = branch.len() * HASH_CYCLE_LEN;
        let mut trace = TraceTable::new(TRACE_WIDTH, trace_length);

        // skip the first node of the branch because it will be computed in the trace as hash(value)
        let branch = &branch[1..];

        trace.fill(
            |state| {
                // initialize first state of the computation
                state[0] = value[0];
                state[1] = value[1];
                state[2..].fill(BaseElement::ZERO);
            },
            |step, state| {
                // execute the transition function for all steps
                //
                // For the first 7 steps of each 8-step cycle, compute a single round of Rescue
                // hash in registers [0..6]. On the 8th step, insert the next branch node into the
                // trace in the positions defined by the next bit of the leaf index. If the bit is
                // ZERO, the next node goes into registers [2, 3], if it is ONE, the node goes into
                // registers [0, 1].

                let cycle_num = step / HASH_CYCLE_LEN;
                let cycle_pos = step % HASH_CYCLE_LEN;

                if cycle_pos < NUM_HASH_ROUNDS {
                    rescue::apply_round(&mut state[..HASH_STATE_WIDTH], step);
                } else {
                    let branch_node = branch[cycle_num].to_elements();
                    let index_bit = BaseElement::new(((index >> cycle_num) & 1) as u128);
                    if index_bit == BaseElement::ZERO {
                        // if index bit is zero, new branch node goes into registers [2, 3]; values
                        // in registers [0, 1] (the accumulated hash) remain unchanged
                        state[2] = branch_node[0];
                        state[3] = branch_node[1];
                    } else {
                        // if index bit is one, accumulated hash goes into registers [2, 3],
                        // and new branch nodes goes into registers [0, 1]
                        state[2] = state[0];
                        state[3] = state[1];
                        state[0] = branch_node[0];
                        state[1] = branch_node[1];
                    }
                    // reset the capacity registers of the state to ZERO
                    state[4] = BaseElement::ZERO;
                    state[5] = BaseElement::ZERO;

                    state[6] = index_bit;
                }
            },
        );

        // set index bit at the second step to one; this still results in a valid execution trace
        // because actual index bits are inserted into the trace after step 7, but it ensures
        // that there are no repeating patterns in the index bit register, and thus the degree
        // of the index bit constraint is stable.
        trace.set(6, 1, FieldElement::ONE);

        trace
    }
}

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

    fn get_pub_inputs(&self, trace: &Self::Trace) -> PublicInputs {
        let last_step = trace.length() - 1;
        PublicInputs {
            tree_root: [trace.get(0, last_step), 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::Blake3_256;
use winterfell::{FieldExtension, ProofOptions};

#[test]
fn merkle_test_basic_proof_verification() {
    let merkle = Box::new(super::MerkleExample::<Blake3_256>::new(
        7,
        build_options(false),
    ));
    crate::tests::test_basic_proof_verification(merkle);
}

#[test]
fn merkle_test_basic_proof_verification_extension() {
    let merkle = Box::new(super::MerkleExample::<Blake3_256>::new(
        7,
        build_options(true),
    ));
    crate::tests::test_basic_proof_verification(merkle);
}

#[test]
fn merkle_test_basic_proof_verification_fail() {
    let merkle = Box::new(super::MerkleExample::<Blake3_256>::new(
        7,
        build_options(false),
    ));
    crate::tests::test_basic_proof_verification_fail(merkle);
}

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


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

use super::{rescue, BaseElement, FieldElement, ProofOptions, CYCLE_LENGTH, TRACE_WIDTH};
use crate::utils::{are_equal, is_zero, not, EvaluationResult};
use winterfell::{
    math::ToElements, Air, AirContext, Assertion, EvaluationFrame, TraceInfo,
    TransitionConstraintDegree,
};

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

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

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

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

impl ToElements<BaseElement> for PublicInputs {
    fn to_elements(&self) -> Vec<BaseElement> {
        let mut result = self.seed.to_vec();
        result.extend_from_slice(&self.result);
        result
    }
}

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

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

    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    fn new(trace_info: TraceInfo, pub_inputs: PublicInputs, options: ProofOptions) -> Self {
        let degrees = vec![
            TransitionConstraintDegree::with_cycles(3, vec![CYCLE_LENGTH]),
            TransitionConstraintDegree::with_cycles(3, vec![CYCLE_LENGTH]),
            TransitionConstraintDegree::with_cycles(3, vec![CYCLE_LENGTH]),
            TransitionConstraintDegree::with_cycles(3, vec![CYCLE_LENGTH]),
        ];
        assert_eq!(TRACE_WIDTH, trace_info.width());
        RescueAir {
            context: AirContext::new(trace_info, degrees, 4, options),
            seed: pub_inputs.seed,
            result: pub_inputs.result,
        }
    }

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

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

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

        // when hash_flag = 1, constraints for Rescue round are enforced
        rescue::enforce_round(result, current, next, ark, hash_flag);

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

    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
        // Assert starting and ending values of the hash chain
        let last_step = self.trace_length() - 1;
        vec![
            Assertion::single(0, 0, self.seed[0]),
            Assertion::single(1, 0, self.seed[1]),
            Assertion::single(0, last_step, self.result[0]),
            Assertion::single(1, last_step, self.result[1]),
        ]
    }

    fn get_periodic_column_values(&self) -> Vec<Vec<Self::BaseField>> {
        let mut result = vec![CYCLE_MASK.to_vec()];
        result.append(&mut rescue::get_round_constants());
        result
    }
}

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

/// when flag = 1, enforces that the next state of the computation is defined like so:
/// - the first two registers are equal to the values from the previous step
/// - the other two registers are equal to 0
fn enforce_hash_copy<E: FieldElement>(result: &mut [E], current: &[E], next: &[E], flag: E) {
    result.agg_constraint(0, flag, are_equal(current[0], next[0]));
    result.agg_constraint(1, flag, are_equal(current[1], next[1]));
    result.agg_constraint(2, flag, is_zero(next[2]));
    result.agg_constraint(3, flag, is_zero(next[3]));
}


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

use 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,
};

#[allow(clippy::module_inception)]
pub(crate) mod rescue;

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

mod prover;
use prover::RescueProver;

#[cfg(test)]
mod tests;

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

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

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

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

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

pub struct RescueExample<H: ElementHasher> {
    options: ProofOptions,
    chain_length: usize,
    seed: [BaseElement; 2],
    result: [BaseElement; 2],
    _hasher: PhantomData<H>,
}

impl<H: ElementHasher> RescueExample<H> {
    pub fn new(chain_length: usize, options: ProofOptions) -> Self {
        assert!(
            chain_length.is_power_of_two(),
            "chain length must a power of 2"
        );
        let seed = [BaseElement::from(42u8), BaseElement::from(43u8)];

        // compute the sequence of hashes using external implementation of Rescue hash
        let now = Instant::now();
        let result = compute_hash_chain(seed, chain_length);
        debug!(
            "Computed a chain of {} Rescue hashes in {} ms",
            chain_length,
            now.elapsed().as_millis(),
        );

        RescueExample {
            options,
            chain_length,
            seed,
            result,
            _hasher: PhantomData,
        }
    }
}

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

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

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

        // generate the execution trace
        let now = Instant::now();
        let trace = prover.build_trace(self.seed, self.chain_length);
        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 {
            seed: self.seed,
            result: self.result,
        };
        winterfell::verify::<RescueAir, H, DefaultRandomCoin<H>>(proof, pub_inputs)
    }

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

// HELPER FUNCTIONS
// ================================================================================================
fn compute_hash_chain(seed: [BaseElement; 2], length: usize) -> [BaseElement; 2] {
    let mut values = seed;
    let mut result = [BaseElement::ZERO; 2];
    for _ in 0..length {
        rescue::hash(values, &mut result);
        values.copy_from_slice(&result);
    }
    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::{
    rescue, BaseElement, DefaultRandomCoin, ElementHasher, FieldElement, PhantomData, ProofOptions,
    Prover, PublicInputs, RescueAir, Trace, TraceTable, CYCLE_LENGTH, NUM_HASH_ROUNDS,
};

// RESCUE PROVER
// ================================================================================================

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

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

    pub fn build_trace(
        &self,
        seed: [BaseElement; 2],
        iterations: usize,
    ) -> TraceTable<BaseElement> {
        // allocate memory to hold the trace table
        let trace_length = iterations * CYCLE_LENGTH;
        let mut trace = TraceTable::new(4, trace_length);

        trace.fill(
            |state| {
                // initialize first state of the computation
                state[0] = seed[0];
                state[1] = seed[1];
                state[2] = BaseElement::ZERO;
                state[3] = BaseElement::ZERO;
            },
            |step, state| {
                // execute the transition function for all steps
                //
                // for the first 14 steps in every cycle, compute a single round of
                // Rescue hash; for the remaining 2 rounds, just carry over the values
                // in the first two registers to the next step
                if (step % CYCLE_LENGTH) < NUM_HASH_ROUNDS {
                    rescue::apply_round(state, step);
                } else {
                    state[2] = BaseElement::ZERO;
                    state[3] = BaseElement::ZERO;
                }
            },
        );

        trace
    }
}

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

    fn get_pub_inputs(&self, trace: &Self::Trace) -> PublicInputs {
        let last_step = trace.length() - 1;
        PublicInputs {
            seed: [trace.get(0, 0), trace.get(1, 0)],
            result: [trace.get(0, last_step), 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 crate::utils::{are_equal, EvaluationResult};
use winterfell::math::{fields::f128::BaseElement, FieldElement};

/// The number of rounds is set to 14 to provide 128-bit security level.
/// computed using algorithm 7 from https://eprint.iacr.org/2020/1143.pdf
const NUM_ROUNDS: usize = 14;

pub const STATE_WIDTH: usize = 4;
const CYCLE_LENGTH: usize = 16;

// HASH FUNCTION
// ================================================================================================

/// Implementation of Rescue hash function with a 4 element state and 14 rounds. Accepts a
/// 2-element input, and returns a 2-element digest.
pub fn hash(value: [BaseElement; 2], result: &mut [BaseElement]) {
    let mut state = BaseElement::zeroed_vector(STATE_WIDTH);
    state[..2].copy_from_slice(&value);
    for i in 0..NUM_ROUNDS {
        apply_round(&mut state, i);
    }
    result.copy_from_slice(&state[..2]);
}

// TRACE
// ================================================================================================

pub fn apply_round(state: &mut [BaseElement], step: usize) {
    // determine which round constants to use
    let ark = ARK[step % CYCLE_LENGTH];

    // apply first half of Rescue round
    apply_sbox(state);
    apply_mds(state);
    add_constants(state, &ark, 0);

    // apply second half of Rescue round
    apply_inv_sbox(state);
    apply_mds(state);
    add_constants(state, &ark, STATE_WIDTH);
}

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

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

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

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

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

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

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

    constants
}

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

#[inline(always)]
#[allow(clippy::needless_range_loop)]
fn add_constants(state: &mut [BaseElement], ark: &[BaseElement], offset: usize) {
    for i in 0..STATE_WIDTH {
        state[i] += ark[offset + i];
    }
}

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

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

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

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

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

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

// RESCUE CONSTANTS
// ================================================================================================
const ALPHA: u32 = 3;
const INV_ALPHA: u128 = 226854911280625642308916371969163307691;

const MDS: [BaseElement; STATE_WIDTH * STATE_WIDTH] = [
    BaseElement::new(340282366920938463463374557953744960808),
    BaseElement::new(1080),
    BaseElement::new(340282366920938463463374557953744961147),
    BaseElement::new(40),
    BaseElement::new(340282366920938463463374557953744932377),
    BaseElement::new(42471),
    BaseElement::new(340282366920938463463374557953744947017),
    BaseElement::new(1210),
    BaseElement::new(340282366920938463463374557953744079447),
    BaseElement::new(1277640),
    BaseElement::new(340282366920938463463374557953744532108),
    BaseElement::new(33880),
    BaseElement::new(340282366920938463463374557953720263017),
    BaseElement::new(35708310),
    BaseElement::new(340282366920938463463374557953733025977),
    BaseElement::new(925771),
];

const INV_MDS: [BaseElement; STATE_WIDTH * STATE_WIDTH] = [
    BaseElement::new(18020639985667067681479625318803400939),
    BaseElement::new(119196285838491236328880430704594968577),
    BaseElement::new(231409255903369280423951003551679307334),
    BaseElement::new(311938552114349342492438056332412246225),
    BaseElement::new(245698978747161380010236204726851770228),
    BaseElement::new(32113671753878130773768090116517402309),
    BaseElement::new(284248318938217584166130208504515171073),
    BaseElement::new(118503764402619831976614612559605579465),
    BaseElement::new(42476948408512208745085164298752800413),
    BaseElement::new(283594571303717652525183978492772054516),
    BaseElement::new(94047455979774690913009073579656179991),
    BaseElement::new(260445758149872374743470899536308888155),
    BaseElement::new(12603050626701424572717576220509072651),
    BaseElement::new(250660673575506110946271793719013778251),
    BaseElement::new(113894235293153614657151429548304212092),
    BaseElement::new(303406774346515776750608316419662860081),
];

pub const ARK: [[BaseElement; STATE_WIDTH * 2]; CYCLE_LENGTH] = [
    [
        BaseElement::new(252629594110556276281235816992330349983),
        BaseElement::new(121163867507455621442731872354015891839),
        BaseElement::new(244623479936175870778515556108748234900),
        BaseElement::new(181999122442017949289616572388308120964),
        BaseElement::new(130035663054758320517176088024859935575),
        BaseElement::new(274932696133623013607933255959111946013),
        BaseElement::new(130096286077538976127585373664362805864),
        BaseElement::new(209506446014122131232133742654202790201),
    ],
    [
        BaseElement::new(51912929769931267810162308005565017268),
        BaseElement::new(202610584823002946089528994694473145326),
        BaseElement::new(295992101426532309592836871256175669136),
        BaseElement::new(313404555247438968545340310449654540090),
        BaseElement::new(137671644572045862038757754124537020379),
        BaseElement::new(29113322527929260506148183779738829778),
        BaseElement::new(98634637270536166954048957710629281939),
        BaseElement::new(90484051915535813802492401077197602516),
    ],
    [
        BaseElement::new(193753019093186599897082621380539177732),
        BaseElement::new(88328997664086495053801384396180288832),
        BaseElement::new(134379598544046716907663161480793367313),
        BaseElement::new(50911186425769400405474055284903795891),
        BaseElement::new(12945394282446072785093894845750344239),
        BaseElement::new(110650301505380365788620562912149942995),
        BaseElement::new(154214463184362737046953674082326221874),
        BaseElement::new(306646039504788072647764955304698381135),
    ],
    [
        BaseElement::new(279745705918489041552127329708931301079),
        BaseElement::new(111293612078035530300709391234153848359),
        BaseElement::new(18110020378502034462498434861690576309),
        BaseElement::new(41797883582559360517115865611622162330),
        BaseElement::new(333888808893608021579859508112201825908),
        BaseElement::new(291192643991850989562610634125476905625),
        BaseElement::new(115042354025120848770557866862388897952),
        BaseElement::new(281483497320099569269754505499721335457),
    ],
    [
        BaseElement::new(172898111753678285350206449646444309824),
        BaseElement::new(202661860135906394577472615378659980424),
        BaseElement::new(141885268042225970011312316000526746741),
        BaseElement::new(270195331267041521741794476882482499817),
        BaseElement::new(196457080224171120865903216527675657315),
        BaseElement::new(56730777565482395039564396246195716949),
        BaseElement::new(4886253806084919544862202000090732791),
        BaseElement::new(147384194551383352824518757380733021990),
    ],
    [
        BaseElement::new(119476237236248181092343711369608370324),
        BaseElement::new(182869361251406039022577235058473348729),
        BaseElement::new(45308522364899994411952744852450066909),
        BaseElement::new(15438528253368638146901598290564135576),
        BaseElement::new(130060283207960095436997328133261743365),
        BaseElement::new(83953475955438079154228277940680487556),
        BaseElement::new(328659226769709797512044291035930357326),
        BaseElement::new(228749522131871685132212950281473676382),
    ],
    [
        BaseElement::new(46194972462682851176957413491161426658),
        BaseElement::new(296333983305826854863835978241833143471),
        BaseElement::new(138957733159616849361016139528307260698),
        BaseElement::new(67842086763518777676559492559456199109),
        BaseElement::new(45580040156133202522383315452912604930),
        BaseElement::new(67567837934606680937620346425373752595),
        BaseElement::new(202860989528104560171546683198384659325),
        BaseElement::new(22630500510153322451285114937258973361),
    ],
    [
        BaseElement::new(324160761097464842200838878419866223614),
        BaseElement::new(338466547889555546143667391979278153877),
        BaseElement::new(189171173535649401433078628567098769571),
        BaseElement::new(162173266902020502126600904559755837464),
        BaseElement::new(136209703129442038834374731074825683052),
        BaseElement::new(61998071517031804812562190829480056772),
        BaseElement::new(307309080039351604461536918194634835054),
        BaseElement::new(26708622949278137915061761772299784349),
    ],
    [
        BaseElement::new(129516553661717764361826568456881002617),
        BaseElement::new(224023580754958002183324313900177991825),
        BaseElement::new(17590440203644538688189654586240082513),
        BaseElement::new(135610063062379124269847491297867667710),
        BaseElement::new(146865534517067293442442506551295645352),
        BaseElement::new(238139104484181583196227119098779158429),
        BaseElement::new(39300761479713744892853256947725570060),
        BaseElement::new(54114440355764484955231402374312070440),
    ],
    [
        BaseElement::new(222758070305343916663075833184045878425),
        BaseElement::new(323840793618712078836672915700599856701),
        BaseElement::new(103586087979277053032666296091805459741),
        BaseElement::new(160263698024385270625527195046420579470),
        BaseElement::new(76620453913654705501329735586535761337),
        BaseElement::new(117793948142462197480091377165008040465),
        BaseElement::new(86998218841589258723143213495722487114),
        BaseElement::new(203188618662906890442620821687773659689),
    ],
    [
        BaseElement::new(313098786815741054633864043424353402357),
        BaseElement::new(133085673687338880872979866135939079867),
        BaseElement::new(219888424885634764555580944265544343421),
        BaseElement::new(5893221169005427793512575133564978746),
        BaseElement::new(123830602624063632344313821515642988189),
        BaseElement::new(99030942908036387138287682010525589136),
        BaseElement::new(181549003357535890945363082242256699137),
        BaseElement::new(152424978799328476472358562493335008209),
    ],
    [
        BaseElement::new(274481943862544603168725464029979191673),
        BaseElement::new(4975004592976331754728718693838357226),
        BaseElement::new(101850445399221640701542169338886750079),
        BaseElement::new(230325699922192981509673754024218912397),
        BaseElement::new(50419227750575087142720761582056939006),
        BaseElement::new(112444234528764731925178653200320603078),
        BaseElement::new(312169855609816651638877239277948636598),
        BaseElement::new(204255114617024487729019111502542629940),
    ],
    [
        BaseElement::new(95797476952346525817251811755749179939),
        BaseElement::new(306977388944722094681694167558392710189),
        BaseElement::new(300754874465668732709232449646112602172),
        BaseElement::new(25567836410351071106804347269705784680),
        BaseElement::new(129659188855548935155840545784705385753),
        BaseElement::new(228441586459539470069565041053012869566),
        BaseElement::new(178382533299631576605259357906020320778),
        BaseElement::new(274458637266680353971597477639962034316),
    ],
    [
        BaseElement::new(280059913840028448065185235205261648486),
        BaseElement::new(246537412674731137211182698562269717969),
        BaseElement::new(259930078572522349821084822750913159564),
        BaseElement::new(186061633995391650657311511040160727356),
        BaseElement::new(179777566992900315528995607912777709520),
        BaseElement::new(209753365793154515863736129686836743468),
        BaseElement::new(270445008049478596978645420017585428243),
        BaseElement::new(70998387591825316724846035292940615733),
    ],
    [BaseElement::ZERO; 8],
    [BaseElement::ZERO; 8],
];


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

use super::Blake3_256;
use winterfell::{FieldExtension, ProofOptions};

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

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

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

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