// 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::mds::mds_f64_8x8::mds_multiply;
use super::{exp_acc, Digest, ElementHasher, Hasher};
use core::convert::TryInto;
use core::ops::Range;
use math::{fields::f64::BaseElement, FieldElement, StarkField};

mod digest;
pub use digest::ElementDigest;

#[cfg(test)]
mod tests;

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

/// Sponge state is set to 8 field elements or 64 bytes; 4 elements are reserved for rate and
/// the remaining 4 elements are reserved for capacity.
const STATE_WIDTH: usize = 8;

/// The rate portion of the state is located in elements 4 through 7.
const RATE_RANGE: Range<usize> = 4..8;
const RATE_WIDTH: usize = RATE_RANGE.end - RATE_RANGE.start;

/// Jive compression mode doesn't consider rate and capacity registers.
const INPUT1_RANGE: Range<usize> = 0..4;
const INPUT2_RANGE: Range<usize> = 4..8;

/// The capacity portion of the state is located in elements 0, 1, 2, and 3.
const CAPACITY_RANGE: Range<usize> = 0..4;

/// The output of the hash function is a digest which consists of 4 field elements or 32 bytes.
///
/// The digest is returned from state elements 4, 5, 6, and 7 (the first four elements of the
/// rate portion).
const DIGEST_RANGE: Range<usize> = 4..8;
const DIGEST_SIZE: usize = DIGEST_RANGE.end - DIGEST_RANGE.start;

/// The number of rounds is set to 7 to target 128-bit security level with 40% security margin;
/// computed using algorithm 7 from <https://eprint.iacr.org/2020/1143.pdf>
const NUM_ROUNDS: usize = 7;

/// S-Box and Inverse S-Box powers;
/// computed using algorithm 6 from <https://eprint.iacr.org/2020/1143.pdf>
///
/// The constants are defined for tests only because the exponentiations in the code are unrolled
/// for efficiency reasons.
#[cfg(test)]
const ALPHA: u64 = 7;
#[cfg(test)]
const INV_ALPHA: u64 = 10540996611094048183;

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

/// Implementation of [Hasher] trait for Rescue Prime hash function with 256-bit output.
///
/// The hash function is implemented according to the Rescue Prime
/// [specifications](https://eprint.iacr.org/2020/1143.pdf) with the following exception:
/// * We set the number of rounds to 7, which implies a 40% security margin instead of the 50%
///   margin used in the specifications (a 50% margin rounds up to 8 rounds). The primary
///   motivation for this is that having the number of rounds be one less than a power of two
///   simplifies AIR design for computations involving the hash function.
/// * When hashing a sequence of elements, implement the Hirose padding rule. However, it also
///   means that our instantiation of Rescue Prime cannot be used in a stream mode as the number
///   of elements to be hashed must be known upfront.
/// * We use the first 4 elements of the state (rather than the last 4 elements of the state) for
///   capacity and the remaining 8 elements for rate. The output of the hash function comes from
///   the first four elements of the rate portion of the state (elements 4, 5, 6, and 7). This
///   effectively applies a fixed bit permutation before and after XLIX permutation. We assert
///   without proof that this does not affect security of the construction.
/// * Instead of using Vandermonde matrices as a standard way of generating an MDS matrix as
///   described in Rescue Prime paper, we use a methodology developed by Polygon Zero to find an
///   MDS matrix with coefficients which are small powers of two in frequency domain. This allows
///   us to dramatically reduce MDS matrix multiplication time. Using a different MDS matrix does
///   not affect security of the hash function as any MDS matrix satisfies Rescue Prime
///   construction (as described in section 4.2 of the paper).
///
/// The parameters used to instantiate the function are:
/// * Field: 64-bit prime field with modulus 2^64 - 2^32 + 1.
/// * State width: 8 field elements.
/// * Capacity size: 4 field elements.
/// * Number of founds: 7.
/// * S-Box degree: 7.
///
/// The above parameters target 128-bit security level. The digest consists of four field elements
/// and it can be serialized into 32 bytes (256 bits).
///
/// ## Hash output consistency
/// Functions [hash_elements()](RpJive64_256::hash_elements), [merge()](RpJive64_256::merge), and
/// [merge_with_int()](RpJive64_256::merge_with_int) are not consistent. This is because the former
/// is instantiated with a sponge construction, while the latter use the Jive compression mode and
/// hence do not rely on the sponge construction.
///
/// In addition, [hash()](RpJive64_256::hash) function is not consistent with the functions mentioned
/// above. For example, if we take two field elements, serialize them to bytes and hash them using
/// [hash()](RpJive64_256::hash), the result will differ from the result obtained by hashing these
/// elements directly using [hash_elements()](RpJive64_256::hash_elements) function. The reason for
/// this difference is that [hash()](RpJive64_256::hash) function needs to be able to handle
/// arbitrary binary strings, which may or may not encode valid field elements - and thus,
/// deserialization procedure used by this function is different from the procedure used to
/// deserialize valid field elements.
///
/// Thus, if the underlying data consists of valid field elements, it might make more sense
/// to deserialize them into field elements and then hash them using
/// [hash_elements()](RpJive64_256::hash_elements) function rather then hashing the serialized bytes
/// using [hash()](RpJive64_256::hash) function.
pub struct RpJive64_256();

impl Hasher for RpJive64_256 {
    type Digest = ElementDigest;

    const COLLISION_RESISTANCE: u32 = 128;

    fn hash(bytes: &[u8]) -> Self::Digest {
        // compute the number of elements required to represent the string; we will be processing
        // the string in 7-byte chunks, thus the number of elements will be equal to the number
        // of such chunks (including a potential partial chunk at the end).
        let num_elements = if bytes.len() % 7 == 0 {
            bytes.len() / 7
        } else {
            bytes.len() / 7 + 1
        };

        // initialize state to all zeros, except for the first element of the capacity part, which
        // is set to 1 if the number of elements is not a multiple of RATE_WIDTH.
        let mut state = [BaseElement::ZERO; STATE_WIDTH];
        if num_elements % RATE_WIDTH != 0 {
            state[CAPACITY_RANGE.start] = BaseElement::ONE;
        }

        // break the string into 7-byte chunks, convert each chunk into a field element, and
        // absorb the element into the rate portion of the state. we use 7-byte chunks because
        // every 7-byte chunk is guaranteed to map to some field element.
        let mut i = 0;
        let mut buf = [0_u8; 8];
        for (index, chunk) in bytes.chunks(7).enumerate() {
            if index < num_elements - 1 {
                buf[..7].copy_from_slice(chunk);
            } else {
                // if we are dealing with the last chunk, it may be smaller than 7 bytes long, so
                // we need to handle it slightly differently. we also append a byte with value 1
                // to the end of the string; this pads the string in such a way that adding
                // trailing zeros results in different hash
                let chunk_len = chunk.len();
                buf = [0_u8; 8];
                buf[..chunk_len].copy_from_slice(chunk);
                buf[chunk_len] = 1;
            }

            // convert the bytes into a field element and absorb it into the rate portion of the
            // state; if the rate is filled up, apply the Griffin permutation and start absorbing
            // again from zero index.
            state[RATE_RANGE.start + i] += BaseElement::new(u64::from_le_bytes(buf));
            i += 1;
            if i % RATE_WIDTH == 0 {
                Self::apply_permutation(&mut state);
                i = 0;
            }
        }

        // if we absorbed some elements but didn't apply a permutation to them (would happen when
        // the number of elements is not a multiple of RATE_WIDTH), apply a final permutation after
        // padding by appending a 1 followed by as many 0 as necessary to make the input length a
        // multiple of the RATE_WIDTH.
        if i > 0 {
            state[RATE_RANGE.start + i] = BaseElement::ONE;
            i += 1;
            while i != RATE_WIDTH {
                state[RATE_RANGE.start + i] = BaseElement::ZERO;
                i += 1;
            }
            Self::apply_permutation(&mut state);
        }

        // return the first 4 elements of the state as hash result
        ElementDigest::new(state[DIGEST_RANGE].try_into().unwrap())
    }

    // We do not rely on the sponge construction to build our compression function. Instead, we use
    // the Jive compression mode designed in https://eprint.iacr.org/2022/840.pdf.
    fn merge(values: &[Self::Digest; 2]) -> Self::Digest {
        // initialize the state by copying the digest elements into the state
        let initial_state: [BaseElement; STATE_WIDTH] = Self::Digest::digests_as_elements(values)
            .try_into()
            .unwrap();
        let mut state = initial_state;

        // apply the Rescue permutation and apply the final Jive summation
        Self::apply_permutation(&mut state);

        Self::apply_jive_summation(&initial_state, &state)
    }

    // We do not rely on the sponge construction to build our compression function. Instead, we use
    // the Jive compression mode designed in https://eprint.iacr.org/2022/840.pdf.
    fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
        // initialize the state as follows:
        // - seed is copied into the first 4 elements of the state.
        // - if the value fits into a single field element, copy it into the fifth rate element
        //   and set the last state element to 5 (the number of elements to be hashed).
        // - if the value doesn't fit into a single field element, split it into two field
        //   elements, copy them into state elements 5 and 6, and set the last state element
        //   to 6.
        let mut state = [BaseElement::ZERO; STATE_WIDTH];
        state[INPUT1_RANGE].copy_from_slice(seed.as_elements());
        state[INPUT2_RANGE.start] = BaseElement::new(value);
        if value < BaseElement::MODULUS {
            state[INPUT2_RANGE.end - 1] = BaseElement::new(DIGEST_SIZE as u64 + 1);
        } else {
            state[INPUT2_RANGE.start + 1] = BaseElement::new(value / BaseElement::MODULUS);
            state[INPUT2_RANGE.end - 1] = BaseElement::new(DIGEST_SIZE as u64 + 2);
        }

        let initial_state = state;
        // apply the Rescue permutation and apply the final Jive summation
        Self::apply_permutation(&mut state);

        Self::apply_jive_summation(&initial_state, &state)
    }
}

impl ElementHasher for RpJive64_256 {
    type BaseField = BaseElement;

    fn hash_elements<E: FieldElement<BaseField = Self::BaseField>>(elements: &[E]) -> Self::Digest {
        // convert the elements into a list of base field elements
        let elements = E::slice_as_base_elements(elements);

        // initialize state to all zeros, except for the first element of the capacity part, which
        // is set to 1 if the number of elements is not a multiple of RATE_WIDTH.
        let mut state = [BaseElement::ZERO; STATE_WIDTH];
        if elements.len() % RATE_WIDTH != 0 {
            state[CAPACITY_RANGE.start] = BaseElement::ONE;
        }

        // absorb elements into the state one by one until the rate portion of the state is filled
        // up; then apply the Griffin permutation and start absorbing again; repeat until all
        // elements have been absorbed
        let mut i = 0;
        for &element in elements.iter() {
            state[RATE_RANGE.start + i] += element;
            i += 1;
            if i % RATE_WIDTH == 0 {
                Self::apply_permutation(&mut state);
                i = 0;
            }
        }

        // if we absorbed some elements but didn't apply a permutation to them (would happen when
        // the number of elements is not a multiple of RATE_WIDTH), apply a final permutation after
        // padding by appending a 1 followed by as many 0 as necessary to make the input length a
        // multiple of the RATE_WIDTH.
        if i > 0 {
            state[RATE_RANGE.start + i] = BaseElement::ONE;
            i += 1;
            while i != RATE_WIDTH {
                state[RATE_RANGE.start + i] = BaseElement::ZERO;
                i += 1;
            }
            Self::apply_permutation(&mut state);
        }

        // return the first 4 elements of the state as hash result
        ElementDigest::new(state[DIGEST_RANGE].try_into().unwrap())
    }
}

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

impl RpJive64_256 {
    // CONSTANTS
    // --------------------------------------------------------------------------------------------

    /// The number of rounds is set to 7 to target 128-bit security level with 40% security margin.
    pub const NUM_ROUNDS: usize = NUM_ROUNDS;

    /// Sponge state is set to 8 field elements or 64 bytes; 4 elements are reserved for rate and
    /// the remaining 4 elements are reserved for capacity.
    pub const STATE_WIDTH: usize = STATE_WIDTH;

    /// The rate portion of the state is located in elements 4 through 7 (inclusive).
    pub const RATE_RANGE: Range<usize> = RATE_RANGE;

    /// The capacity portion of the state is located in elements 0, 1, 2, and 3.
    pub const CAPACITY_RANGE: Range<usize> = CAPACITY_RANGE;

    /// The output of the hash function can be read from state elements 4, 5, 6, and 7.
    pub const DIGEST_RANGE: Range<usize> = DIGEST_RANGE;

    /// MDS matrix used for computing the linear layer in a Rescue Prime round.
    pub const MDS: [[BaseElement; STATE_WIDTH]; STATE_WIDTH] = MDS;

    /// Inverse of the MDS matrix.
    pub const INV_MDS: [[BaseElement; STATE_WIDTH]; STATE_WIDTH] = INV_MDS;

    /// Round constants added to the hasher state in the first half of the Rescue Prime round.
    pub const ARK1: [[BaseElement; STATE_WIDTH]; NUM_ROUNDS] = ARK1;

    /// Round constants added to the hasher state in the second half of the Rescue Prime round.
    pub const ARK2: [[BaseElement; STATE_WIDTH]; NUM_ROUNDS] = ARK2;

    // RESCUE PERMUTATION
    // --------------------------------------------------------------------------------------------

    /// Applies Rescue-XLIX permutation to the provided state.
    pub fn apply_permutation(state: &mut [BaseElement; STATE_WIDTH]) {
        // implementation is based on algorithm 3 from <https://eprint.iacr.org/2020/1143.pdf>
        // apply round function 7 times; this provides 128-bit security with 40% security margin
        for i in 0..NUM_ROUNDS {
            Self::apply_round(state, i);
        }
    }

    /// Rescue-XLIX round function.
    #[inline(always)]
    pub fn apply_round(state: &mut [BaseElement; STATE_WIDTH], round: usize) {
        // apply first half of Rescue round
        Self::apply_sbox(state);
        Self::apply_mds(state);
        Self::add_constants(state, &ARK1[round]);

        // apply second half of Rescue round
        Self::apply_inv_sbox(state);
        Self::apply_mds(state);
        Self::add_constants(state, &ARK2[round]);
    }

    #[inline(always)]
    pub fn apply_jive_summation(
        initial_state: &[BaseElement; STATE_WIDTH],
        final_state: &[BaseElement; STATE_WIDTH],
    ) -> ElementDigest {
        let mut result = [BaseElement::ZERO; DIGEST_SIZE];
        for (i, r) in result.iter_mut().enumerate() {
            *r = initial_state[i]
                + initial_state[DIGEST_SIZE + i]
                + final_state[i]
                + final_state[DIGEST_SIZE + i];
        }

        ElementDigest::new(result)
    }

    // HELPER FUNCTIONS
    // --------------------------------------------------------------------------------------------

    #[inline(always)]
    fn apply_mds(state: &mut [BaseElement; STATE_WIDTH]) {
        mds_multiply(state)
    }

    #[inline(always)]
    fn add_constants(state: &mut [BaseElement; STATE_WIDTH], ark: &[BaseElement; STATE_WIDTH]) {
        state.iter_mut().zip(ark).for_each(|(s, &k)| *s += k);
    }

    #[inline(always)]
    fn apply_sbox(state: &mut [BaseElement; STATE_WIDTH]) {
        state[0] = state[0].exp7();
        state[1] = state[1].exp7();
        state[2] = state[2].exp7();
        state[3] = state[3].exp7();
        state[4] = state[4].exp7();
        state[5] = state[5].exp7();
        state[6] = state[6].exp7();
        state[7] = state[7].exp7();
    }

    #[inline(always)]
    fn apply_inv_sbox(state: &mut [BaseElement; STATE_WIDTH]) {
        // compute base^10540996611094048183 using 72 multiplications per array element
        // 10540996611094048183 = b1001001001001001001001001001000110110110110110110110110110110111

        // compute base^10
        let mut t1 = *state;
        t1.iter_mut().for_each(|t| *t = t.square());

        // compute base^100
        let mut t2 = t1;
        t2.iter_mut().for_each(|t| *t = t.square());

        // compute base^100100
        let t3 = exp_acc::<BaseElement, STATE_WIDTH, 3>(t2, t2);

        // compute base^100100100100
        let t4 = exp_acc::<BaseElement, STATE_WIDTH, 6>(t3, t3);

        // compute base^100100100100100100100100
        let t5 = exp_acc::<BaseElement, STATE_WIDTH, 12>(t4, t4);

        // compute base^100100100100100100100100100100
        let t6 = exp_acc::<BaseElement, STATE_WIDTH, 6>(t5, t3);

        // compute base^1001001001001001001001001001000100100100100100100100100100100
        let t7 = exp_acc::<BaseElement, STATE_WIDTH, 31>(t6, t6);

        // compute base^1001001001001001001001001001000110110110110110110110110110110111
        for (i, s) in state.iter_mut().enumerate() {
            let a = (t7[i].square() * t6[i]).square().square();
            let b = t1[i] * t2[i] * *s;
            *s = a * b;
        }
    }
}

// MDS
// ================================================================================================
/// Rescue MDS matrix
const MDS: [[BaseElement; STATE_WIDTH]; STATE_WIDTH] = [
    [
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(21),
        BaseElement::new(8),
    ],
    [
        BaseElement::new(8),
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(21),
    ],
    [
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(7),
        BaseElement::new(6),
    ],
    [
        BaseElement::new(6),
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(7),
    ],
    [
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(13),
        BaseElement::new(10),
    ],
    [
        BaseElement::new(10),
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(13),
    ],
    [
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(23),
        BaseElement::new(8),
    ],
    [
        BaseElement::new(8),
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(23),
    ],
];

/// Rescue Inverse MDS matrix
const INV_MDS: [[BaseElement; STATE_WIDTH]; STATE_WIDTH] = [
    [
        BaseElement::new(10671399028204489528),
        BaseElement::new(15436289366139187412),
        BaseElement::new(4624329233769728317),
        BaseElement::new(18200084821960740316),
        BaseElement::new(8736112961492104393),
        BaseElement::new(1953609990965186349),
        BaseElement::new(12477339747250042564),
        BaseElement::new(1495657543820456485),
    ],
    [
        BaseElement::new(1495657543820456485),
        BaseElement::new(10671399028204489528),
        BaseElement::new(15436289366139187412),
        BaseElement::new(4624329233769728317),
        BaseElement::new(18200084821960740316),
        BaseElement::new(8736112961492104393),
        BaseElement::new(1953609990965186349),
        BaseElement::new(12477339747250042564),
    ],
    [
        BaseElement::new(12477339747250042564),
        BaseElement::new(1495657543820456485),
        BaseElement::new(10671399028204489528),
        BaseElement::new(15436289366139187412),
        BaseElement::new(4624329233769728317),
        BaseElement::new(18200084821960740316),
        BaseElement::new(8736112961492104393),
        BaseElement::new(1953609990965186349),
    ],
    [
        BaseElement::new(1953609990965186349),
        BaseElement::new(12477339747250042564),
        BaseElement::new(1495657543820456485),
        BaseElement::new(10671399028204489528),
        BaseElement::new(15436289366139187412),
        BaseElement::new(4624329233769728317),
        BaseElement::new(18200084821960740316),
        BaseElement::new(8736112961492104393),
    ],
    [
        BaseElement::new(8736112961492104393),
        BaseElement::new(1953609990965186349),
        BaseElement::new(12477339747250042564),
        BaseElement::new(1495657543820456485),
        BaseElement::new(10671399028204489528),
        BaseElement::new(15436289366139187412),
        BaseElement::new(4624329233769728317),
        BaseElement::new(18200084821960740316),
    ],
    [
        BaseElement::new(18200084821960740316),
        BaseElement::new(8736112961492104393),
        BaseElement::new(1953609990965186349),
        BaseElement::new(12477339747250042564),
        BaseElement::new(1495657543820456485),
        BaseElement::new(10671399028204489528),
        BaseElement::new(15436289366139187412),
        BaseElement::new(4624329233769728317),
    ],
    [
        BaseElement::new(4624329233769728317),
        BaseElement::new(18200084821960740316),
        BaseElement::new(8736112961492104393),
        BaseElement::new(1953609990965186349),
        BaseElement::new(12477339747250042564),
        BaseElement::new(1495657543820456485),
        BaseElement::new(10671399028204489528),
        BaseElement::new(15436289366139187412),
    ],
    [
        BaseElement::new(15436289366139187412),
        BaseElement::new(4624329233769728317),
        BaseElement::new(18200084821960740316),
        BaseElement::new(8736112961492104393),
        BaseElement::new(1953609990965186349),
        BaseElement::new(12477339747250042564),
        BaseElement::new(1495657543820456485),
        BaseElement::new(10671399028204489528),
    ],
];

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

/// Rescue round constants;
/// computed using algorithm 5 from <https://eprint.iacr.org/2020/1143.pdf>
///
/// The constants are broken up into two arrays ARK1 and ARK2; ARK1 contains the constants for the
/// first half of Rescue round, and ARK2 contains constants for the second half of Rescue round.
const ARK1: [[BaseElement; STATE_WIDTH]; NUM_ROUNDS] = [
    [
        BaseElement::new(5250156239823432273),
        BaseElement::new(17991370199276831394),
        BaseElement::new(15363758995121189373),
        BaseElement::new(7550390719632034712),
        BaseElement::new(705744964663370588),
        BaseElement::new(14718080047998507086),
        BaseElement::new(15612952641514293932),
        BaseElement::new(8827614218997241655),
    ],
    [
        BaseElement::new(3820104553051581938),
        BaseElement::new(3385456123263281593),
        BaseElement::new(16094709323995557719),
        BaseElement::new(16303336019291352506),
        BaseElement::new(8678496957982514796),
        BaseElement::new(498270172890916765),
        BaseElement::new(17676155962043649331),
        BaseElement::new(14993644560894569061),
    ],
    [
        BaseElement::new(14258773164148374760),
        BaseElement::new(1655972393090532756),
        BaseElement::new(7105012644980738960),
        BaseElement::new(11852376844296856307),
        BaseElement::new(17816158174482938174),
        BaseElement::new(3981864273667206359),
        BaseElement::new(2807469273751819673),
        BaseElement::new(14974221859211617968),
    ],
    [
        BaseElement::new(15947271309323471269),
        BaseElement::new(14698197888879866148),
        BaseElement::new(14077077040726269118),
        BaseElement::new(2859805440338816615),
        BaseElement::new(4945184696648790387),
        BaseElement::new(15183288803792940883),
        BaseElement::new(7601775560447886378),
        BaseElement::new(6477224812816853098),
    ],
    [
        BaseElement::new(18213733347601447845),
        BaseElement::new(10031679943792626621),
        BaseElement::new(5971928707867502549),
        BaseElement::new(4916840084933933812),
        BaseElement::new(3613815642787339926),
        BaseElement::new(16715066477165606893),
        BaseElement::new(14603075385258290966),
        BaseElement::new(6037771699330759024),
    ],
    [
        BaseElement::new(11092469678405138663),
        BaseElement::new(14512788091784891767),
        BaseElement::new(12690682422447262976),
        BaseElement::new(4807355108863118656),
        BaseElement::new(5207405791308193025),
        BaseElement::new(5970889292753030887),
        BaseElement::new(17691092604759176390),
        BaseElement::new(2731892623388788619),
    ],
    [
        BaseElement::new(9320990164295747317),
        BaseElement::new(8313044787501051613),
        BaseElement::new(15388579942433649113),
        BaseElement::new(16827303822369113172),
        BaseElement::new(7362247368635881413),
        BaseElement::new(5501558211335089067),
        BaseElement::new(16959364163466644433),
        BaseElement::new(15127897185888596873),
    ],
];

const ARK2: [[BaseElement; STATE_WIDTH]; NUM_ROUNDS] = [
    [
        BaseElement::new(9860068499471230379),
        BaseElement::new(10391494434594667033),
        BaseElement::new(4986587677027284267),
        BaseElement::new(17781977240739864050),
        BaseElement::new(6888921375142581299),
        BaseElement::new(8950831725295674725),
        BaseElement::new(17048848277806802259),
        BaseElement::new(14146306451370933851),
    ],
    [
        BaseElement::new(707569561928852298),
        BaseElement::new(6724851263229096394),
        BaseElement::new(16052786826295129381),
        BaseElement::new(1966016718617096590),
        BaseElement::new(9416027981257317341),
        BaseElement::new(650995073054283087),
        BaseElement::new(10013853213448688130),
        BaseElement::new(14400137552134409897),
    ],
    [
        BaseElement::new(7149263702162640230),
        BaseElement::new(7096225564191267298),
        BaseElement::new(12197502430442379401),
        BaseElement::new(12804378092281676880),
        BaseElement::new(17409570408925731570),
        BaseElement::new(2819914464281065415),
        BaseElement::new(15831648359524824910),
        BaseElement::new(15629743966484525526),
    ],
    [
        BaseElement::new(17953398529773387863),
        BaseElement::new(6198711330432012203),
        BaseElement::new(9157726872360640492),
        BaseElement::new(9493333679697066249),
        BaseElement::new(16030612341681265024),
        BaseElement::new(4739709630031417239),
        BaseElement::new(18287301685877696586),
        BaseElement::new(8798230489526342293),
    ],
    [
        BaseElement::new(11624786627634502148),
        BaseElement::new(12924370583547723043),
        BaseElement::new(11192385058160295505),
        BaseElement::new(14350900531623057057),
        BaseElement::new(6649040255431543914),
        BaseElement::new(2106567763792008889),
        BaseElement::new(12434281915569617273),
        BaseElement::new(8101377239551798417),
    ],
    [
        BaseElement::new(13925815041351874730),
        BaseElement::new(15981136477777934021),
        BaseElement::new(17398194123970783302),
        BaseElement::new(17377636820017036987),
        BaseElement::new(5173992930377549692),
        BaseElement::new(3688194845376511083),
        BaseElement::new(16177005022792194790),
        BaseElement::new(6482787365501773067),
    ],
    [
        BaseElement::new(9197066592623932055),
        BaseElement::new(1777435748159421921),
        BaseElement::new(5079482957444239813),
        BaseElement::new(15080163201683705054),
        BaseElement::new(4278835591662809119),
        BaseElement::new(6609842793229774583),
        BaseElement::new(651644751771720476),
        BaseElement::new(14434199410773467460),
    ],
];


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

use super::{
    BaseElement, ElementDigest, ElementHasher, FieldElement, Hasher, RpJive64_256, StarkField,
    ALPHA, INV_ALPHA, INV_MDS, MDS, STATE_WIDTH,
};
use core::convert::TryInto;
use proptest::prelude::*;

use rand_utils::{rand_array, rand_value};

#[test]
fn mds_inv_test() {
    let mut mul_result = [[BaseElement::new(0); STATE_WIDTH]; STATE_WIDTH];
    for i in 0..STATE_WIDTH {
        for j in 0..STATE_WIDTH {
            let result = {
                let mut result = BaseElement::new(0);
                for k in 0..STATE_WIDTH {
                    result += MDS[i][k] * INV_MDS[k][j]
                }
                result
            };
            mul_result[i][j] = result;
            if i == j {
                assert_eq!(result, BaseElement::new(1));
            } else {
                assert_eq!(result, BaseElement::new(0));
            }
        }
    }
}
#[test]
fn test_alphas() {
    let e: BaseElement = rand_value();
    let e_exp = e.exp(ALPHA.into());
    assert_eq!(e, e_exp.exp(INV_ALPHA));
}

#[test]
fn test_sbox() {
    let state: [BaseElement; STATE_WIDTH] = rand_array();

    let mut expected = state;
    expected.iter_mut().for_each(|v| *v = v.exp(ALPHA));

    let mut actual = state;
    RpJive64_256::apply_sbox(&mut actual);

    assert_eq!(expected, actual);
}

#[test]
fn test_inv_sbox() {
    let state: [BaseElement; STATE_WIDTH] = rand_array();

    let mut expected = state;
    expected.iter_mut().for_each(|v| *v = v.exp(INV_ALPHA));

    let mut actual = state;
    RpJive64_256::apply_inv_sbox(&mut actual);

    assert_eq!(expected, actual);
}

#[test]
fn apply_permutation() {
    let mut state: [BaseElement; STATE_WIDTH] = [
        BaseElement::new(0),
        BaseElement::new(1),
        BaseElement::new(2),
        BaseElement::new(3),
        BaseElement::new(4),
        BaseElement::new(5),
        BaseElement::new(6),
        BaseElement::new(7),
    ];

    RpJive64_256::apply_permutation(&mut state);

    // expected values are obtained by executing sage reference implementation code
    let expected = vec![
        BaseElement::new(16940713730596720799),
        BaseElement::new(16218555904323712189),
        BaseElement::new(11042680722444601138),
        BaseElement::new(5370396747047489939),
        BaseElement::new(6349480890410006944),
        BaseElement::new(1551053614279730715),
        BaseElement::new(3995941143622927528),
        BaseElement::new(9350074312471431779),
    ];

    assert_eq!(expected, state);
}

#[test]
fn hash_elements_vs_merge() {
    let elements: [BaseElement; 8] = rand_array();

    let digests: [ElementDigest; 2] = [
        ElementDigest::new(elements[..4].try_into().unwrap()),
        ElementDigest::new(elements[4..].try_into().unwrap()),
    ];

    let m_result = RpJive64_256::merge(&digests);
    let h_result = RpJive64_256::hash_elements(&elements);

    // Because we use the Jive compression mode, `merge` and
    // `hash_elements` methods are incompatible.
    assert_ne!(m_result, h_result);
}

#[test]
fn hash_elements_vs_merge_with_int() {
    let seed = ElementDigest::new(rand_array());

    // ----- value fits into a field element ------------------------------------------------------
    let val: BaseElement = rand_value();
    let m_result = RpJive64_256::merge_with_int(seed, val.as_int());

    let mut elements = seed.as_elements().to_vec();
    elements.push(val);
    let h_result = RpJive64_256::hash_elements(&elements);

    // Because we use the Jive compression mode, `merge` and
    // `hash_elements` methods are incompatible.
    assert_ne!(m_result, h_result);

    // ----- value does not fit into a field element ----------------------------------------------
    let val = BaseElement::MODULUS + 2;
    let m_result = RpJive64_256::merge_with_int(seed, val);

    let mut elements = seed.as_elements().to_vec();
    elements.push(BaseElement::new(val));
    elements.push(BaseElement::new(1));
    let h_result = RpJive64_256::hash_elements(&elements);

    // Because we use the Jive compression mode, `merge` and
    // `hash_elements` methods are incompatible.
    assert_ne!(m_result, h_result);
}

#[test]
fn hash_padding() {
    // adding a zero bytes at the end of a byte string should result in a different hash
    let r1 = RpJive64_256::hash(&[1_u8, 2, 3]);
    let r2 = RpJive64_256::hash(&[1_u8, 2, 3, 0]);
    assert_ne!(r1, r2);

    // same as above but with bigger inputs
    let r1 = RpJive64_256::hash(&[1_u8, 2, 3, 4, 5, 6]);
    let r2 = RpJive64_256::hash(&[1_u8, 2, 3, 4, 5, 6, 0]);
    assert_ne!(r1, r2);

    // same as above but with input splitting over two elements
    let r1 = RpJive64_256::hash(&[1_u8, 2, 3, 4, 5, 6, 7]);
    let r2 = RpJive64_256::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0]);
    assert_ne!(r1, r2);

    // same as above but with multiple zeros
    let r1 = RpJive64_256::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0, 0]);
    let r2 = RpJive64_256::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0, 0, 0, 0]);
    assert_ne!(r1, r2);
}

#[test]
fn hash_elements_padding() {
    let e1: [BaseElement; 2] = rand_array();
    let e2 = [e1[0], e1[1], BaseElement::ZERO];

    let r1 = RpJive64_256::hash_elements(&e1);
    let r2 = RpJive64_256::hash_elements(&e2);
    assert_ne!(r1, r2);
}

#[inline(always)]
fn apply_mds_naive(state: &mut [BaseElement; STATE_WIDTH]) {
    let mut result = [BaseElement::ZERO; STATE_WIDTH];
    result.iter_mut().zip(MDS).for_each(|(r, mds_row)| {
        state.iter().zip(mds_row).for_each(|(&s, m)| {
            *r += m * s;
        });
    });
    *state = result;
}

proptest! {
    #[test]
    fn mds_freq_proptest(a in any::<[u64;STATE_WIDTH]>()) {

        let mut v1 = [BaseElement::ZERO;STATE_WIDTH];
        let mut v2;

        for i in 0..STATE_WIDTH {
            v1[i] = BaseElement::new(a[i]);
        }
        v2 = v1.clone();

        apply_mds_naive(&mut v1);
        RpJive64_256::apply_mds(&mut v2);

        prop_assert_eq!(v1, v2);
    }
}


// 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::{ByteDigest, ElementHasher, Hasher};
use core::marker::PhantomData;
use math::{FieldElement, StarkField};
use sha3::Digest;
use utils::ByteWriter;

// SHA3 WITH 256-BIT OUTPUT
// ================================================================================================

/// Implementation of the [Hasher](super::Hasher) trait for SHA3 hash function with 256-bit
/// output.
pub struct Sha3_256<B: StarkField>(PhantomData<B>);

impl<B: StarkField> Hasher for Sha3_256<B> {
    type Digest = ByteDigest<32>;

    const COLLISION_RESISTANCE: u32 = 128;

    fn hash(bytes: &[u8]) -> Self::Digest {
        ByteDigest(sha3::Sha3_256::digest(bytes).into())
    }

    fn merge(values: &[Self::Digest; 2]) -> Self::Digest {
        ByteDigest(sha3::Sha3_256::digest(ByteDigest::digests_as_bytes(values)).into())
    }

    fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
        let mut data = [0; 40];
        data[..32].copy_from_slice(&seed.0);
        data[32..].copy_from_slice(&value.to_le_bytes());
        ByteDigest(sha3::Sha3_256::digest(data).into())
    }
}

impl<B: StarkField> ElementHasher for Sha3_256<B> {
    type BaseField = B;

    fn hash_elements<E: FieldElement<BaseField = Self::BaseField>>(elements: &[E]) -> Self::Digest {
        if B::IS_CANONICAL {
            // when element's internal and canonical representations are the same, we can hash
            // element bytes directly
            let bytes = E::elements_as_bytes(elements);
            ByteDigest(sha3::Sha3_256::digest(bytes).into())
        } else {
            // when elements' internal and canonical representations differ, we need to serialize
            // them before hashing
            let mut hasher = ShaHasher::new();
            hasher.write(elements);
            ByteDigest(hasher.finalize())
        }
    }
}

// SHA HASHER
// ================================================================================================

/// Wrapper around SHA3 hasher to implement [ByteWriter] trait for it.
struct ShaHasher(sha3::Sha3_256);

impl ShaHasher {
    pub fn new() -> Self {
        Self(sha3::Sha3_256::new())
    }

    pub fn finalize(self) -> [u8; 32] {
        self.0.finalize().into()
    }
}

impl ByteWriter for ShaHasher {
    fn write_u8(&mut self, value: u8) {
        self.0.update([value]);
    }

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


// 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::Hasher;
use core::slice;
use utils::{collections::Vec, iterators::*, rayon};

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

pub const MIN_CONCURRENT_LEAVES: usize = 1024;

// PUBLIC FUNCTIONS
// ================================================================================================

/// Builds all internal nodes of the Merkle using all available threads and stores the
/// results in a single vector such that root of the tree is at position 1, nodes immediately
/// under the root is at positions 2 and 3 etc.
pub fn build_merkle_nodes<H: Hasher>(leaves: &[H::Digest]) -> Vec<H::Digest> {
    let n = leaves.len() / 2;

    // create un-initialized array to hold all intermediate nodes
    let mut nodes = unsafe { utils::uninit_vector::<H::Digest>(2 * n) };
    nodes[0] = H::Digest::default();

    // re-interpret leaves as an array of two leaves fused together and use it to
    // build first row of internal nodes (parents of leaves)
    let two_leaves = unsafe { slice::from_raw_parts(leaves.as_ptr() as *const [H::Digest; 2], n) };
    nodes[n..]
        .par_iter_mut()
        .zip(two_leaves.par_iter())
        .for_each(|(target, source)| *target = H::merge(source));

    // calculate all other tree nodes, we can't use regular iterators  here because
    // access patterns are rather complicated - so, we use regular threads instead

    // number of sub-trees must always be a power of 2
    let num_subtrees = rayon::current_num_threads().next_power_of_two();
    let batch_size = n / num_subtrees;

    // re-interpret nodes as an array of two nodes fused together
    let two_nodes = unsafe { slice::from_raw_parts(nodes.as_ptr() as *const [H::Digest; 2], n) };

    // process each subtree in a separate thread
    rayon::scope(|s| {
        for i in 0..num_subtrees {
            let nodes = unsafe { &mut *(&mut nodes[..] as *mut [H::Digest]) };
            s.spawn(move |_| {
                let mut batch_size = batch_size / 2;
                let mut start_idx = n / 2 + batch_size * i;
                while start_idx >= num_subtrees {
                    for k in (start_idx..(start_idx + batch_size)).rev() {
                        nodes[k] = H::merge(&two_nodes[k]);
                    }
                    start_idx /= 2;
                    batch_size /= 2;
                }
            });
        }
    });

    // finish the tip of the tree
    for i in (1..num_subtrees).rev() {
        nodes[i] = H::merge(&two_nodes[i]);
    }

    nodes
}

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

#[cfg(test)]
mod tests {
    use crate::hash::{ByteDigest, Sha3_256};
    use math::fields::f128::BaseElement;
    use proptest::collection::vec;
    use proptest::prelude::*;

    proptest! {
        #[test]
        fn build_merkle_nodes_concurrent(ref data in vec(any::<[u8; 32]>(), 256..257).no_shrink()) {
            let leaves = ByteDigest::bytes_as_digests(&data).to_vec();
            let sequential = super::super::build_merkle_nodes::<Sha3_256<BaseElement>>(&leaves);
            let concurrent = super::build_merkle_nodes::<Sha3_256<BaseElement>>(&leaves);
            assert_eq!(concurrent, sequential);
        }
    }
}


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

use crate::{errors::MerkleTreeError, hash::Hasher};
use core::slice;
use utils::collections::{BTreeMap, BTreeSet, Vec};

mod proofs;
pub use proofs::BatchMerkleProof;

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

#[cfg(test)]
mod tests;

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

/// A fully-balanced Merkle tree.
///
/// In this implementation, a Merkle tree consists of two types of nodes: leaves and internal nodes
/// (one of which is a tree root). All nodes must be instances of the digest specified by the
/// [Hasher] used to build the tree.
///
/// ```text
///       *        <- tree root
///     /   \
///    /     \
///   *       *    <- internal nodes
///  / \     / \
/// o   o   o   o  <- leaves
/// |   |   |   |
/// #   #   #   #  <- values
/// ```
///
/// A tree can be built from a slice of leaves using [MerkleTree::new()] function. Thus, the user
/// is responsible for performing the first level of hashing (i.e., hashing values into leaf
/// nodes). The number of leaves must always be a power of two so that the tree is fully balanced,
/// and a tree must contain at least two leaves.
///
/// The depth of a tree is zero-based. Thus, a tree with two leaves has depth 1, a tree with four
/// leaves has depth 2 etc.
///
/// When the crate is compiled with `concurrent` feature enabled, tree construction will be
/// performed in multiple threads (usually, as many threads as there are logical cores on the
/// machine). The number of threads can be configured via `RAYON_NUM_THREADS` environment variable.
///
/// To generate an inclusion proof for a given leaf, [MerkleTree::prove()] method can be used.
/// You can also use [MerkleTree::prove_batch()] method to generate inclusion proofs for multiple
/// leaves. The advantage of the batch method is that redundant internal nodes are removed from
/// the batch proof, thereby compressing it (we use a variation of the
/// [Octopus](https://eprint.iacr.org/2017/933) algorithm).
///
/// To verify proofs, [MerkleTree::verify()] and [MerkleTree::verify_batch()] functions can be
/// used respectively.
///
/// # Examples
/// ```
/// # use winter_crypto::{MerkleTree, Hasher, hashers::Blake3_256};
/// # use math::fields::f128::BaseElement;
/// type Blake3 = Blake3_256::<BaseElement>;
///
/// // build a tree
/// let leaves = [
///     Blake3::hash(&[1u8]),
///     Blake3::hash(&[2u8]),
///     Blake3::hash(&[3u8]),
///     Blake3::hash(&[4u8]),
/// ];
/// let tree = MerkleTree::<Blake3>::new(leaves.to_vec()).unwrap();
/// assert_eq!(2, tree.depth());
/// assert_eq!(leaves, tree.leaves());
///
/// // generate a proof
/// let proof = tree.prove(2).unwrap();
/// assert_eq!(3, proof.len());
/// assert_eq!(leaves[2], proof[0]);
///
/// // verify proof
/// assert!(MerkleTree::<Blake3>::verify(*tree.root(), 2, &proof).is_ok());
/// assert!(MerkleTree::<Blake3>::verify(*tree.root(), 1, &proof).is_err());
/// ```
#[derive(Debug)]
pub struct MerkleTree<H: Hasher> {
    nodes: Vec<H::Digest>,
    leaves: Vec<H::Digest>,
}

// MERKLE TREE IMPLEMENTATION
// ================================================================================================

impl<H: Hasher> MerkleTree<H> {
    // CONSTRUCTORS
    // --------------------------------------------------------------------------------------------

    /// Returns new Merkle tree built from the provide leaves using hash function specified by the
    /// `H` generic parameter.
    ///
    /// When `concurrent` feature is enabled, the tree is built using multiple threads.
    ///
    /// # Errors
    /// Returns an error if:
    /// * Fewer than two leaves were provided.
    /// * Number of leaves is not a power of two.
    pub fn new(leaves: Vec<H::Digest>) -> Result<Self, MerkleTreeError> {
        if leaves.len() < 2 {
            return Err(MerkleTreeError::TooFewLeaves(2, leaves.len()));
        }
        if !leaves.len().is_power_of_two() {
            return Err(MerkleTreeError::NumberOfLeavesNotPowerOfTwo(leaves.len()));
        }

        #[cfg(not(feature = "concurrent"))]
        let nodes = build_merkle_nodes::<H>(&leaves);

        #[cfg(feature = "concurrent")]
        let nodes = if leaves.len() <= concurrent::MIN_CONCURRENT_LEAVES {
            build_merkle_nodes::<H>(&leaves)
        } else {
            concurrent::build_merkle_nodes::<H>(&leaves)
        };

        Ok(MerkleTree { nodes, leaves })
    }

    /// Forms a MerkleTree from a list of nodes and leaves.
    ///
    /// Nodes are supplied as a vector where the root is stored at position 1.
    ///
    /// # Errors
    /// Returns an error if:
    /// * Fewer than two leaves were provided.
    /// * Number of leaves is not a power of two.
    ///
    /// # Panics
    /// Panics if nodes doesn't have the same length as leaves.
    pub fn from_raw_parts(
        nodes: Vec<H::Digest>,
        leaves: Vec<H::Digest>,
    ) -> Result<Self, MerkleTreeError> {
        if leaves.len() < 2 {
            return Err(MerkleTreeError::TooFewLeaves(2, leaves.len()));
        }
        if !leaves.len().is_power_of_two() {
            return Err(MerkleTreeError::NumberOfLeavesNotPowerOfTwo(leaves.len()));
        }
        assert_eq!(nodes.len(), leaves.len());
        Ok(MerkleTree { nodes, leaves })
    }

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

    /// Returns the root of the tree.
    pub fn root(&self) -> &H::Digest {
        &self.nodes[1]
    }

    /// Returns depth of the tree.
    ///
    /// The depth of a tree is zero-based. Thus, a tree with two leaves has depth 1, a tree with
    /// four leaves has depth 2 etc.
    pub fn depth(&self) -> usize {
        self.leaves.len().ilog2() as usize
    }

    /// Returns leaf nodes of the tree.
    pub fn leaves(&self) -> &[H::Digest] {
        &self.leaves
    }

    // PROVING METHODS
    // --------------------------------------------------------------------------------------------

    /// Returns a Merkle path to a leaf at the specified `index`.
    ///
    /// The leaf itself will be the first element in the path.
    ///
    /// # Errors
    /// Returns an error if the specified index is greater than or equal to the number of leaves
    /// in the tree.
    pub fn prove(&self, index: usize) -> Result<Vec<H::Digest>, MerkleTreeError> {
        if index >= self.leaves.len() {
            return Err(MerkleTreeError::LeafIndexOutOfBounds(
                self.leaves.len(),
                index,
            ));
        }

        let mut proof = vec![self.leaves[index], self.leaves[index ^ 1]];

        let mut index = (index + self.nodes.len()) >> 1;
        while index > 1 {
            proof.push(self.nodes[index ^ 1]);
            index >>= 1;
        }

        Ok(proof)
    }

    /// Computes Merkle paths for the provided indexes and compresses the paths into a single proof.
    ///
    /// # Errors
    /// Returns an error if:
    /// * No indexes were provided (i.e., `indexes` is an empty slice).
    /// * Number of provided indexes is greater than 255.
    /// * Any of the provided indexes are greater than or equal to the number of leaves in the
    ///   tree.
    /// * List of indexes contains duplicates.
    pub fn prove_batch(&self, indexes: &[usize]) -> Result<BatchMerkleProof<H>, MerkleTreeError> {
        if indexes.is_empty() {
            return Err(MerkleTreeError::TooFewLeafIndexes);
        }
        if indexes.len() > proofs::MAX_PATHS {
            return Err(MerkleTreeError::TooManyLeafIndexes(
                proofs::MAX_PATHS,
                indexes.len(),
            ));
        }

        let index_map = map_indexes(indexes, self.depth())?;
        let indexes = normalize_indexes(indexes);
        let mut leaves = vec![H::Digest::default(); index_map.len()];
        let mut nodes: Vec<Vec<H::Digest>> = Vec::with_capacity(indexes.len());

        // populate the proof with leaf node values
        let n = self.leaves.len();
        let mut next_indexes: Vec<usize> = Vec::new();
        for index in indexes {
            let missing: Vec<H::Digest> = (index..index + 2)
                .flat_map(|i| {
                    let v = self.leaves[i];
                    if let Some(idx) = index_map.get(&i) {
                        leaves[*idx] = v;
                        None
                    } else {
                        Some(v)
                    }
                })
                .collect();
            nodes.push(missing);

            next_indexes.push((index + n) >> 1);
        }

        // add required internal nodes to the proof, skipping redundancies
        for _ in 1..self.depth() {
            let indexes = next_indexes.clone();
            next_indexes.truncate(0);

            let mut i = 0;
            while i < indexes.len() {
                let sibling_index = indexes[i] ^ 1;
                if i + 1 < indexes.len() && indexes[i + 1] == sibling_index {
                    i += 1;
                } else {
                    nodes[i].push(self.nodes[sibling_index]);
                }

                // add parent index to the set of next indexes
                next_indexes.push(sibling_index >> 1);

                i += 1;
            }
        }

        Ok(BatchMerkleProof {
            leaves,
            nodes,
            depth: self.depth() as u8,
        })
    }

    // VERIFICATION METHODS
    // --------------------------------------------------------------------------------------------

    /// Checks whether the `proof` for the specified `index` is valid.
    ///
    /// # Errors
    /// Returns an error if the specified `proof` (which is a Merkle path) does not resolve to the
    /// specified `root`.
    pub fn verify(
        root: H::Digest,
        index: usize,
        proof: &[H::Digest],
    ) -> Result<(), MerkleTreeError> {
        let r = index & 1;
        let mut v = H::merge(&[proof[r], proof[1 - r]]);

        let mut index = (index + 2usize.pow((proof.len() - 1) as u32)) >> 1;
        for &p in proof.iter().skip(2) {
            v = if index & 1 == 0 {
                H::merge(&[v, p])
            } else {
                H::merge(&[p, v])
            };
            index >>= 1;
        }

        if v != root {
            return Err(MerkleTreeError::InvalidProof);
        }
        Ok(())
    }

    /// Checks whether the batch proof contains Merkle paths for the of the specified `indexes`.
    ///
    /// # Errors
    /// Returns an error if:
    /// * No indexes were provided (i.e., `indexes` is an empty slice).
    /// * Number of provided indexes is greater than 255.
    /// * Any of the specified `indexes` is greater than or equal to the number of leaves in the
    ///   tree from which the batch proof was generated.
    /// * List of indexes contains duplicates.
    /// * Any of the paths in the batch proof does not resolve to the specified `root`.
    pub fn verify_batch(
        root: &H::Digest,
        indexes: &[usize],
        proof: &BatchMerkleProof<H>,
    ) -> Result<(), MerkleTreeError> {
        if *root != proof.get_root(indexes)? {
            return Err(MerkleTreeError::InvalidProof);
        }
        Ok(())
    }
}

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

/// Returns the internal nodes of a Merkle tree defined by the specified leaves.
///
/// The internal nodes are turned as a vector where the root is stored at position 1, its children
/// are stored at positions 2, 3, their children are stored at positions 4, 5, 6, 7 etc.
///
/// This function is exposed primarily for benchmarking purposes. It is not intended to be used
/// directly by the end users of the crate.
pub fn build_merkle_nodes<H: Hasher>(leaves: &[H::Digest]) -> Vec<H::Digest> {
    let n = leaves.len() / 2;

    // create un-initialized array to hold all intermediate nodes
    let mut nodes = unsafe { utils::uninit_vector::<H::Digest>(2 * n) };
    nodes[0] = H::Digest::default();

    // re-interpret leaves as an array of two leaves fused together
    let two_leaves = unsafe { slice::from_raw_parts(leaves.as_ptr() as *const [H::Digest; 2], n) };

    // build first row of internal nodes (parents of leaves)
    for (i, j) in (0..n).zip(n..nodes.len()) {
        nodes[j] = H::merge(&two_leaves[i]);
    }

    // re-interpret nodes as an array of two nodes fused together
    let two_nodes = unsafe { slice::from_raw_parts(nodes.as_ptr() as *const [H::Digest; 2], n) };

    // calculate all other tree nodes
    for i in (1..n).rev() {
        nodes[i] = H::merge(&two_nodes[i]);
    }

    nodes
}

fn map_indexes(
    indexes: &[usize],
    tree_depth: usize,
) -> Result<BTreeMap<usize, usize>, MerkleTreeError> {
    let num_leaves = 2usize.pow(tree_depth as u32);
    let mut map = BTreeMap::new();
    for (i, index) in indexes.iter().cloned().enumerate() {
        map.insert(index, i);
        if index >= num_leaves {
            return Err(MerkleTreeError::LeafIndexOutOfBounds(num_leaves, index));
        }
    }

    if indexes.len() != map.len() {
        return Err(MerkleTreeError::DuplicateLeafIndex);
    }

    Ok(map)
}

fn normalize_indexes(indexes: &[usize]) -> Vec<usize> {
    let mut set = BTreeSet::new();
    for &index in indexes {
        set.insert(index - (index & 1));
    }
    set.into_iter().collect()
}


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

use crate::{errors::MerkleTreeError, Hasher};
use utils::{
    collections::{BTreeMap, Vec},
    string::ToString,
    ByteReader, Deserializable, DeserializationError, Serializable,
};

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

pub(super) const MAX_PATHS: usize = 255;

// BATCH MERKLE PROOF
// ================================================================================================

/// Multiple Merkle paths aggregated into a single proof.
///
/// The aggregation is done in a way which removes all duplicate internal nodes, and thus,
/// it is possible to achieve non-negligible compression as compared to naively concatenating
/// individual Merkle paths. The algorithm is for aggregation is a variation of
/// [Octopus](https://eprint.iacr.org/2017/933).
///
/// Currently, at most 255 paths can be aggregated into a single proof. This limitation is
/// imposed primarily for serialization purposes.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BatchMerkleProof<H: Hasher> {
    /// The leaves being proven
    pub leaves: Vec<H::Digest>,
    /// Hashes of Merkle Tree proof values above the leaf layer
    pub nodes: Vec<Vec<H::Digest>>,
    /// Depth of the leaves
    pub depth: u8,
}

impl<H: Hasher> BatchMerkleProof<H> {
    /// Constructs a batch Merkle proof from individual Merkle authentication paths.
    ///
    /// # Panics
    /// Panics if:
    /// * No paths have been provided (i.e., `paths` is an empty slice).
    /// * More than 255 paths have been provided.
    /// * Number of paths is not equal to the number of indexes.
    /// * Not all paths have the same length.
    pub fn from_paths(paths: &[Vec<H::Digest>], indexes: &[usize]) -> BatchMerkleProof<H> {
        // TODO: optimize this to reduce amount of vector cloning.
        assert!(!paths.is_empty(), "at least one path must be provided");
        assert!(
            paths.len() <= MAX_PATHS,
            "number of paths cannot exceed {MAX_PATHS}"
        );
        assert_eq!(
            paths.len(),
            indexes.len(),
            "number of paths must equal number of indexes"
        );

        let depth = paths[0].len();

        // sort indexes in ascending order, and also re-arrange paths accordingly
        let mut path_map = BTreeMap::new();
        for (&index, path) in indexes.iter().zip(paths.iter().cloned()) {
            assert_eq!(depth, path.len(), "not all paths have the same length");
            path_map.insert(index, path);
        }
        let indexes = path_map.keys().cloned().collect::<Vec<_>>();
        let paths = path_map.values().cloned().collect::<Vec<_>>();
        path_map.clear();

        let mut leaves = vec![H::Digest::default(); indexes.len()];
        let mut nodes: Vec<Vec<H::Digest>> = Vec::with_capacity(indexes.len());

        // populate values and the first layer of proof nodes
        let mut i = 0;
        while i < indexes.len() {
            leaves[i] = paths[i][0];
            if indexes.len() > i + 1 && are_siblings(indexes[i], indexes[i + 1]) {
                leaves[i + 1] = paths[i][1];
                nodes.push(vec![]);
                i += 1;
            } else {
                nodes.push(vec![paths[i][1]]);
            }
            path_map.insert(indexes[i] >> 1, paths[i].clone());
            i += 1;
        }

        // populate all remaining layers of proof nodes
        for d in 2..depth {
            let indexes = path_map.keys().cloned().collect::<Vec<_>>();
            let mut next_path_map = BTreeMap::new();

            let mut i = 0;
            while i < indexes.len() {
                let index = indexes[i];
                let path = path_map.get(&index).unwrap();
                if indexes.len() > i + 1 && are_siblings(index, indexes[i + 1]) {
                    i += 1;
                } else {
                    nodes[i].push(path[d]);
                }
                next_path_map.insert(index >> 1, path.clone());
                i += 1;
            }

            core::mem::swap(&mut path_map, &mut next_path_map);
        }

        BatchMerkleProof {
            leaves,
            nodes,
            depth: (depth - 1) as u8,
        }
    }

    /// Computes a node to which all Merkle paths aggregated in this proof resolve.
    ///
    /// # Errors
    /// Returns an error if:
    /// * No indexes were provided (i.e., `indexes` is an empty slice).
    /// * Number of provided indexes is greater than 255.
    /// * Any of the specified `indexes` is greater than or equal to the number of leaves in the
    ///   tree for which this batch proof was generated.
    /// * List of indexes contains duplicates.
    /// * The proof does not resolve to a single root.
    pub fn get_root(&self, indexes: &[usize]) -> Result<H::Digest, MerkleTreeError> {
        if indexes.is_empty() {
            return Err(MerkleTreeError::TooFewLeafIndexes);
        }
        if indexes.len() > MAX_PATHS {
            return Err(MerkleTreeError::TooManyLeafIndexes(
                MAX_PATHS,
                indexes.len(),
            ));
        }

        let mut buf = [H::Digest::default(); 2];
        let mut v = BTreeMap::new();

        // replace odd indexes, offset, and sort in ascending order
        let index_map = super::map_indexes(indexes, self.depth as usize)?;
        let indexes = super::normalize_indexes(indexes);
        if indexes.len() != self.nodes.len() {
            return Err(MerkleTreeError::InvalidProof);
        }

        // for each index use values to compute parent nodes
        let offset = 2usize.pow(self.depth as u32);
        let mut next_indexes: Vec<usize> = Vec::new();
        let mut proof_pointers: Vec<usize> = Vec::with_capacity(indexes.len());
        for (i, index) in indexes.into_iter().enumerate() {
            // copy values of leaf sibling leaf nodes into the buffer
            match index_map.get(&index) {
                Some(&index1) => {
                    if self.leaves.len() <= index1 {
                        return Err(MerkleTreeError::InvalidProof);
                    }
                    buf[0] = self.leaves[index1];
                    match index_map.get(&(index + 1)) {
                        Some(&index2) => {
                            if self.leaves.len() <= index2 {
                                return Err(MerkleTreeError::InvalidProof);
                            }
                            buf[1] = self.leaves[index2];
                            proof_pointers.push(0);
                        }
                        None => {
                            if self.nodes[i].is_empty() {
                                return Err(MerkleTreeError::InvalidProof);
                            }
                            buf[1] = self.nodes[i][0];
                            proof_pointers.push(1);
                        }
                    }
                }
                None => {
                    if self.nodes[i].is_empty() {
                        return Err(MerkleTreeError::InvalidProof);
                    }
                    buf[0] = self.nodes[i][0];
                    match index_map.get(&(index + 1)) {
                        Some(&index2) => {
                            if self.leaves.len() <= index2 {
                                return Err(MerkleTreeError::InvalidProof);
                            }
                            buf[1] = self.leaves[index2];
                        }
                        None => return Err(MerkleTreeError::InvalidProof),
                    }
                    proof_pointers.push(1);
                }
            }

            // hash sibling nodes into their parent
            let parent = H::merge(&buf);

            let parent_index = (offset + index) >> 1;
            v.insert(parent_index, parent);
            next_indexes.push(parent_index);
        }

        // iteratively move up, until we get to the root
        for _ in 1..self.depth {
            let indexes = next_indexes.clone();
            next_indexes.truncate(0);

            let mut i = 0;
            while i < indexes.len() {
                let node_index = indexes[i];
                let sibling_index = node_index ^ 1;

                // determine the sibling
                let sibling: H::Digest;
                if i + 1 < indexes.len() && indexes[i + 1] == sibling_index {
                    sibling = match v.get(&sibling_index) {
                        Some(sibling) => *sibling,
                        None => return Err(MerkleTreeError::InvalidProof),
                    };
                    i += 1;
                } else {
                    let pointer = proof_pointers[i];
                    if self.nodes[i].len() <= pointer {
                        return Err(MerkleTreeError::InvalidProof);
                    }
                    sibling = self.nodes[i][pointer];
                    proof_pointers[i] += 1;
                }

                // get the node from the map of hashed nodes
                let node = match v.get(&node_index) {
                    Some(node) => node,
                    None => return Err(MerkleTreeError::InvalidProof),
                };

                // compute parent node from node and sibling
                if node_index & 1 != 0 {
                    buf[0] = sibling;
                    buf[1] = *node;
                } else {
                    buf[0] = *node;
                    buf[1] = sibling;
                }
                let parent = H::merge(&buf);

                // add the parent node to the next set of nodes
                let parent_index = node_index >> 1;
                v.insert(parent_index, parent);
                next_indexes.push(parent_index);

                i += 1;
            }
        }
        v.remove(&1).ok_or(MerkleTreeError::InvalidProof)
    }

    /// Computes the uncompressed Merkle paths which aggregate to this proof.
    ///
    /// # Errors
    /// Returns an error if:
    /// * No indexes were provided (i.e., `indexes` is an empty slice).
    /// * Number of provided indexes is greater than 255.
    /// * Number of provided indexes does not match the number of leaf nodes in the proof.
    pub fn into_paths(self, indexes: &[usize]) -> Result<Vec<Vec<H::Digest>>, MerkleTreeError> {
        if indexes.is_empty() {
            return Err(MerkleTreeError::TooFewLeafIndexes);
        }
        if indexes.len() > MAX_PATHS {
            return Err(MerkleTreeError::TooManyLeafIndexes(
                MAX_PATHS,
                indexes.len(),
            ));
        }
        if indexes.len() != self.leaves.len() {
            return Err(MerkleTreeError::InvalidProof);
        }

        let mut partial_tree_map = BTreeMap::new();

        for (&i, leaf) in indexes.iter().zip(self.leaves.iter()) {
            partial_tree_map.insert(i + (1 << (self.depth)), *leaf);
        }

        let mut buf = [H::Digest::default(); 2];
        let mut v = BTreeMap::new();

        // replace odd indexes, offset, and sort in ascending order
        let original_indexes = indexes;
        let index_map = super::map_indexes(indexes, self.depth as usize)?;
        let indexes = super::normalize_indexes(indexes);
        if indexes.len() != self.nodes.len() {
            return Err(MerkleTreeError::InvalidProof);
        }

        // for each index use values to compute parent nodes
        let offset = 2usize.pow(self.depth as u32);
        let mut next_indexes: Vec<usize> = Vec::new();
        let mut proof_pointers: Vec<usize> = Vec::with_capacity(indexes.len());
        for (i, index) in indexes.into_iter().enumerate() {
            // copy values of leaf sibling leaf nodes into the buffer
            match index_map.get(&index) {
                Some(&index1) => {
                    if self.leaves.len() <= index1 {
                        return Err(MerkleTreeError::InvalidProof);
                    }
                    buf[0] = self.leaves[index1];
                    match index_map.get(&(index + 1)) {
                        Some(&index2) => {
                            if self.leaves.len() <= index2 {
                                return Err(MerkleTreeError::InvalidProof);
                            }
                            buf[1] = self.leaves[index2];
                            proof_pointers.push(0);
                        }
                        None => {
                            if self.nodes[i].is_empty() {
                                return Err(MerkleTreeError::InvalidProof);
                            }
                            buf[1] = self.nodes[i][0];
                            proof_pointers.push(1);
                        }
                    }
                }
                None => {
                    if self.nodes[i].is_empty() {
                        return Err(MerkleTreeError::InvalidProof);
                    }
                    buf[0] = self.nodes[i][0];
                    match index_map.get(&(index + 1)) {
                        Some(&index2) => {
                            if self.leaves.len() <= index2 {
                                return Err(MerkleTreeError::InvalidProof);
                            }
                            buf[1] = self.leaves[index2];
                        }
                        None => return Err(MerkleTreeError::InvalidProof),
                    }
                    proof_pointers.push(1);
                }
            }

            // hash sibling nodes into their parent and add it to partial_tree
            let parent = H::merge(&buf);
            partial_tree_map.insert(offset + index, buf[0]);
            partial_tree_map.insert((offset + index) ^ 1, buf[1]);
            let parent_index = (offset + index) >> 1;
            v.insert(parent_index, parent);
            next_indexes.push(parent_index);
            partial_tree_map.insert(parent_index, parent);
        }

        // iteratively move up, until we get to the root
        for _ in 1..self.depth {
            let indexes = next_indexes.clone();
            next_indexes.clear();

            let mut i = 0;
            while i < indexes.len() {
                let node_index = indexes[i];
                let sibling_index = node_index ^ 1;

                // determine the sibling
                let sibling = if i + 1 < indexes.len() && indexes[i + 1] == sibling_index {
                    i += 1;
                    match v.get(&sibling_index) {
                        Some(sibling) => *sibling,
                        None => return Err(MerkleTreeError::InvalidProof),
                    }
                } else {
                    let pointer = proof_pointers[i];
                    if self.nodes[i].len() <= pointer {
                        return Err(MerkleTreeError::InvalidProof);
                    }
                    proof_pointers[i] += 1;
                    self.nodes[i][pointer]
                };

                // get the node from the map of hashed nodes
                let node = match v.get(&node_index) {
                    Some(node) => node,
                    None => return Err(MerkleTreeError::InvalidProof),
                };

                // compute parent node from node and sibling
                partial_tree_map.insert(node_index ^ 1, sibling);
                let parent = if node_index & 1 != 0 {
                    H::merge(&[sibling, *node])
                } else {
                    H::merge(&[*node, sibling])
                };

                // add the parent node to the next set of nodes and partial_tree
                let parent_index = node_index >> 1;
                v.insert(parent_index, parent);
                next_indexes.push(parent_index);
                partial_tree_map.insert(parent_index, parent);

                i += 1;
            }
        }

        original_indexes
            .iter()
            .map(|&i| get_path::<H>(i, &partial_tree_map, self.depth as usize))
            .collect()
    }

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

    /// Converts all internal proof nodes into a vector of bytes.
    ///
    /// # Panics
    /// Panics if:
    /// * The proof contains more than 255 Merkle paths.
    /// * The Merkle paths consist of more than 255 nodes.
    pub fn serialize_nodes(&self) -> Vec<u8> {
        let mut result = Vec::new();

        // record total number of node vectors
        assert!(self.nodes.len() <= u8::MAX as usize, "too many paths");
        result.push(self.nodes.len() as u8);

        // record each node vector as individual bytes
        for nodes in self.nodes.iter() {
            assert!(nodes.len() <= u8::MAX as usize, "too many nodes");
            // record the number of nodes, and append all nodes to the paths buffer
            result.push(nodes.len() as u8);
            for node in nodes.iter() {
                result.append(&mut node.to_bytes());
            }
        }

        result
    }

    /// Parses internal nodes from the provided `node_bytes`, and constructs a batch Merkle proof
    /// from these nodes, provided `leaves`, and provided tree `depth`.
    ///
    /// # Errors
    /// Returns an error if:
    /// * No leaves were provided (i.e., `leaves` is an empty slice).
    /// * Number of provided leaves is greater than 255.
    /// * Tree `depth` was set to zero.
    /// * `node_bytes` could not be deserialized into a valid set of internal nodes.
    pub fn deserialize<R: ByteReader>(
        node_bytes: &mut R,
        leaves: Vec<H::Digest>,
        depth: u8,
    ) -> Result<Self, DeserializationError> {
        if depth == 0 {
            return Err(DeserializationError::InvalidValue(
                "tree depth must be greater than zero".to_string(),
            ));
        }
        if leaves.is_empty() {
            return Err(DeserializationError::InvalidValue(
                "at lease one leaf must be provided".to_string(),
            ));
        }
        if leaves.len() > MAX_PATHS {
            return Err(DeserializationError::InvalidValue(format!(
                "number of leaves cannot exceed {}, but {} were provided",
                MAX_PATHS,
                leaves.len()
            )));
        }

        let num_node_vectors = node_bytes.read_u8()? as usize;
        let mut nodes = Vec::with_capacity(num_node_vectors);
        for _ in 0..num_node_vectors {
            // read the number of digests in the vector
            let num_digests = node_bytes.read_u8()? as usize;

            // read the digests and add them to the node vector
            let digests = H::Digest::read_batch_from(node_bytes, num_digests)?;
            nodes.push(digests);
        }

        Ok(BatchMerkleProof {
            leaves,
            nodes,
            depth,
        })
    }
}

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

/// Two nodes are siblings if index of the left node is even and right node
/// immediately follows the left node.
fn are_siblings(left: usize, right: usize) -> bool {
    left & 1 == 0 && right - 1 == left
}

/// Computes the Merkle path from the computed (partial) tree.
pub fn get_path<H: Hasher>(
    index: usize,
    tree: &BTreeMap<usize, <H as Hasher>::Digest>,
    depth: usize,
) -> Result<Vec<H::Digest>, MerkleTreeError> {
    let mut index = index + (1 << depth);
    let leaf = if let Some(leaf) = tree.get(&index) {
        *leaf
    } else {
        return Err(MerkleTreeError::InvalidProof);
    };

    let mut proof = vec![leaf];
    while index > 1 {
        let leaf = if let Some(leaf) = tree.get(&(index ^ 1)) {
            *leaf
        } else {
            return Err(MerkleTreeError::InvalidProof);
        };

        proof.push(leaf);
        index >>= 1;
    }

    Ok(proof)
}


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

use crate::{errors::RandomCoinError, Digest, ElementHasher, RandomCoin};
use core::convert::TryInto;
use math::{FieldElement, StarkField};
use utils::collections::Vec;

// DEFAULT RANDOM COIN IMPLEMENTATION
// ================================================================================================

/// Pseudo-random element generator for finite fields, which is a default implementation of the
/// RandomCoin trait.
///
/// A random coin can be used to draw elements uniformly at random from the specified base field
/// or from any extension of the base field.
///
/// Internally we use a cryptographic hash function (which is specified via the `H` type parameter),
/// to draw elements from the field. The coin works roughly as follows:
/// - The internal state of the coin consists of a `seed` and a `counter`. At instantiation
///   time, the `seed` is set to a hash of the provided bytes, and the `counter` is set to 0.
/// - To draw the next element, we increment the `counter` and compute hash(`seed` || `counter`).
///   If the resulting value is a valid field element, we return the result; otherwise we try
///   again until a valid element is found or the number of allowed tries is exceeded.
/// - We can also re-seed the coin with a new value. During the reseeding procedure, the
///   seed is set to hash(`old_seed` || `new_seed`), and the counter is reset to 0.
///
/// # Examples
/// ```
/// # use winter_crypto::{RandomCoin, DefaultRandomCoin, Hasher, hashers::Blake3_256};
/// # use math::fields::f128::BaseElement;
/// // initial elements for seeding the random coin
/// let seed = &[BaseElement::new(1), BaseElement::new(2), BaseElement::new(3), BaseElement::new(4)];
///
/// // instantiate a random coin using BLAKE3 as the hash function
/// let mut coin = DefaultRandomCoin::<Blake3_256<BaseElement>>::new(seed);
///
/// // should draw different elements each time
/// let e1 = coin.draw::<BaseElement>().unwrap();;
/// let e2 = coin.draw::<BaseElement>().unwrap();;
/// assert_ne!(e1, e2);
///
/// let e3 = coin.draw::<BaseElement>().unwrap();;
/// assert_ne!(e1, e3);
/// assert_ne!(e2, e3);
///
/// // should draw same elements for the same seed
/// let mut coin2 = DefaultRandomCoin::<Blake3_256<BaseElement>>::new(seed);
/// let mut coin1 = DefaultRandomCoin::<Blake3_256<BaseElement>>::new(seed);
/// let e1 = coin1.draw::<BaseElement>().unwrap();;
/// let e2 = coin2.draw::<BaseElement>().unwrap();;
/// assert_eq!(e1, e2);
///
/// // should draw different elements based on seed
/// let mut coin1 = DefaultRandomCoin::<Blake3_256<BaseElement>>::new(seed);
/// let seed = &[BaseElement::new(2), BaseElement::new(3), BaseElement::new(4), BaseElement::new(5)];
/// let mut coin2 = DefaultRandomCoin::<Blake3_256<BaseElement>>::new(seed);
/// let e1 = coin1.draw::<BaseElement>().unwrap();;
/// let e2 = coin2.draw::<BaseElement>().unwrap();;
/// assert_ne!(e1, e2);
/// ```
pub struct DefaultRandomCoin<H: ElementHasher> {
    seed: H::Digest,
    counter: u64,
}

impl<H: ElementHasher> DefaultRandomCoin<H> {
    /// Updates the state by incrementing the counter and returns hash(seed || counter)
    fn next(&mut self) -> H::Digest {
        self.counter += 1;
        H::merge_with_int(self.seed, self.counter)
    }
}

impl<B: StarkField, H: ElementHasher<BaseField = B>> RandomCoin for DefaultRandomCoin<H> {
    type BaseField = B;
    type Hasher = H;

    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new random coin instantiated with the provided `seed`.
    fn new(seed: &[Self::BaseField]) -> Self {
        let seed = H::hash_elements(seed);
        Self { seed, counter: 0 }
    }

    // RESEEDING
    // --------------------------------------------------------------------------------------------

    /// Reseeds the coin with the specified data by setting the new seed to hash(`seed` || `data`).
    ///
    /// # Examples
    /// ```
    /// # use winter_crypto::{RandomCoin, DefaultRandomCoin, Hasher, hashers::Blake3_256};
    /// # use math::fields::f128::BaseElement;
    /// // initial elements for seeding the random coin
    /// let seed = &[BaseElement::new(1), BaseElement::new(2), BaseElement::new(3), BaseElement::new(4)];
    ///
    /// let mut coin1 = DefaultRandomCoin::<Blake3_256<BaseElement>>::new(seed);
    /// let mut coin2 = DefaultRandomCoin::<Blake3_256<BaseElement>>::new(seed);
    ///
    /// // should draw the same element form both coins
    /// let e1 = coin1.draw::<BaseElement>().unwrap();
    /// let e2 = coin2.draw::<BaseElement>().unwrap();;
    /// assert_eq!(e1, e2);
    ///
    /// // after reseeding should draw different elements
    /// coin2.reseed(Blake3_256::<BaseElement>::hash(&[2, 3, 4, 5]));
    /// let e1 = coin1.draw::<BaseElement>().unwrap();;
    /// let e2 = coin2.draw::<BaseElement>().unwrap();;
    /// assert_ne!(e1, e2);
    /// ```
    fn reseed(&mut self, data: H::Digest) {
        self.seed = H::merge(&[self.seed, data]);
        self.counter = 0;
    }

    /// Reseeds the coin with the specified value by setting the new seed to hash(`seed` ||
    /// `value`).
    ///
    /// # Examples
    /// ```
    /// # use winter_crypto::{RandomCoin, DefaultRandomCoin, Hasher, hashers::Blake3_256};
    /// # use math::fields::f128::BaseElement;
    /// // initial elements for seeding the random coin
    /// let seed = &[BaseElement::new(1), BaseElement::new(2), BaseElement::new(3), BaseElement::new(4)];
    ///
    /// let mut coin1 = DefaultRandomCoin::<Blake3_256<BaseElement>>::new(seed);
    /// let mut coin2 = DefaultRandomCoin::<Blake3_256<BaseElement>>::new(seed);
    ///
    /// // should draw the same element form both coins
    /// let e1 = coin1.draw::<BaseElement>().unwrap();;
    /// let e2 = coin2.draw::<BaseElement>().unwrap();;
    /// assert_eq!(e1, e2);
    ///
    /// // after reseeding should draw different elements
    /// coin2.reseed_with_int(42);
    /// let e1 = coin1.draw::<BaseElement>().unwrap();;
    /// let e2 = coin2.draw::<BaseElement>().unwrap();;
    /// assert_ne!(e1, e2);
    /// ```
    fn reseed_with_int(&mut self, value: u64) {
        self.seed = H::merge_with_int(self.seed, value);
        self.counter = 0;
    }

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

    /// Returns the number of leading zeros in the seed if it is interpreted as an integer in
    /// big-endian byte order.
    ///
    /// # Examples
    /// ```
    /// # use winter_crypto::{RandomCoin, DefaultRandomCoin, Hasher, hashers::Blake3_256};
    /// # use math::fields::f128::BaseElement;
    /// // initial elements for seeding the random coin
    /// let seed = &[BaseElement::new(1), BaseElement::new(2), BaseElement::new(3), BaseElement::new(4)];
    ///
    /// let mut coin = DefaultRandomCoin::<Blake3_256<BaseElement>>::new(seed);
    ///
    /// let mut value = 0;
    /// while coin.check_leading_zeros(value) < 2 {
    ///     value += 1;
    /// }
    ///
    /// coin.reseed_with_int(value);
    /// assert!(coin.leading_zeros() >= 2);
    /// ```
    fn leading_zeros(&self) -> u32 {
        let bytes = self.seed.as_bytes();
        let seed_head = u64::from_le_bytes(bytes[..8].try_into().unwrap());
        seed_head.trailing_zeros()
    }

    /// Computes hash(`seed` || `value`) and returns the number of leading zeros in the resulting
    /// value if it is interpreted as an integer in big-endian byte order.
    fn check_leading_zeros(&self, value: u64) -> u32 {
        let new_seed = H::merge_with_int(self.seed, value);
        let bytes = new_seed.as_bytes();
        let seed_head = u64::from_le_bytes(bytes[..8].try_into().unwrap());
        seed_head.trailing_zeros()
    }

    // DRAW METHODS
    // --------------------------------------------------------------------------------------------

    /// Returns the next pseudo-random field element.
    ///
    /// # Errors
    /// Returns an error if a valid field element could not be generated after 1000 calls to the
    /// PRNG.
    fn draw<E: FieldElement>(&mut self) -> Result<E, RandomCoinError> {
        for _ in 0..1000 {
            // get the next pseudo-random value and take the first ELEMENT_BYTES from it
            let value = self.next();
            let bytes = &value.as_bytes()[..E::ELEMENT_BYTES];

            // check if the bytes can be converted into a valid field element; if they can,
            // return; otherwise try again
            if let Some(element) = E::from_random_bytes(bytes) {
                return Ok(element);
            }
        }

        Err(RandomCoinError::FailedToDrawFieldElement(1000))
    }

    /// Returns a vector of unique integers selected from the range [0, domain_size).
    ///
    /// # Errors
    /// Returns an error if the specified number of unique integers could not be generated
    /// after 1000 calls to the PRNG.
    ///
    /// # Panics
    /// Panics if:
    /// - `domain_size` is not a power of two.
    /// - `num_values` is greater than or equal to `domain_size`.
    ///
    /// # Examples
    /// ```
    /// # use std::collections::HashSet;
    /// # use winter_crypto::{RandomCoin, DefaultRandomCoin, Hasher, hashers::Blake3_256};
    /// # use math::fields::f128::BaseElement;
    /// // initial elements for seeding the random coin
    /// let seed = &[BaseElement::new(1), BaseElement::new(2), BaseElement::new(3), BaseElement::new(4)];
    ///
    /// let mut coin = DefaultRandomCoin::<Blake3_256<BaseElement>>::new(seed);
    ///
    /// let num_values = 20;
    /// let domain_size = 64;
    /// let values = coin.draw_integers(num_values, domain_size).unwrap();
    ///
    /// assert_eq!(num_values, values.len());
    ///
    /// let mut value_set = HashSet::new();
    /// for value in values {
    ///     assert!(value < domain_size);
    ///     assert!(value_set.insert(value));
    /// }
    /// ```
    fn draw_integers(
        &mut self,
        num_values: usize,
        domain_size: usize,
    ) -> Result<Vec<usize>, RandomCoinError> {
        assert!(
            domain_size.is_power_of_two(),
            "domain size must be a power of two"
        );
        assert!(
            num_values < domain_size,
            "number of values must be smaller than domain size"
        );

        // determine how many bits are needed to represent valid values in the domain
        let v_mask = (domain_size - 1) as u64;

        // draw values from PRNG until we get as many unique values as specified by num_queries
        let mut values = Vec::new();
        for _ in 0..1000 {
            // get the next pseudo-random value and read the first 8 bytes from it
            let bytes: [u8; 8] = self.next().as_bytes()[..8].try_into().unwrap();

            // convert to integer and limit the integer to the number of bits which can fit
            // into the specified domain
            let value = (u64::from_le_bytes(bytes) & v_mask) as usize;

            if values.contains(&value) {
                continue;
            }
            values.push(value);
            if values.len() == num_values {
                break;
            }
        }

        if values.len() < num_values {
            return Err(RandomCoinError::FailedToDrawIntegers(
                num_values,
                values.len(),
                1000,
            ));
        }

        Ok(values)
    }
}


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

use crate::{errors::RandomCoinError, ElementHasher, Hasher};
use math::{FieldElement, StarkField};
use utils::collections::Vec;

mod default;
pub use default::DefaultRandomCoin;

// RANDOM COIN TRAIT
// ================================================================================================

/// Pseudo-random element generator for finite fields.
///
/// A random coin can be used to draw elements uniformly at random from the specified base field
/// or from any extension of the base field.
///
/// Internally we use a cryptographic hash function (which is specified via the `Hasher` associated
/// type), to draw elements from the field.
pub trait RandomCoin: Sync {
    /// Base field for random elements which can be generated by this random coin.
    type BaseField: StarkField;

    /// Hash function which is used by the random coin to generate random field elements.
    type Hasher: ElementHasher<BaseField = Self::BaseField>;

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

    /// Returns a new random coin instantiated with the provided `seed`.
    fn new(seed: &[Self::BaseField]) -> Self;

    /// Reseeds the coin with the specified data by setting the new seed to hash(`seed` || `data`).
    fn reseed(&mut self, data: <Self::Hasher as Hasher>::Digest);

    /// Reseeds the coin with the specified value by setting the new seed to hash(`seed` ||
    /// `value`).
    fn reseed_with_int(&mut self, value: u64);

    /// Returns the number of leading zeros in the seed if it is interpreted as an integer in
    /// big-endian byte order.
    fn leading_zeros(&self) -> u32;

    /// Computes hash(`seed` || `value`) and returns the number of leading zeros in the resulting
    /// value if it is interpreted as an integer in big-endian byte order.
    fn check_leading_zeros(&self, value: u64) -> u32;

    /// Returns the next pseudo-random field element.
    ///
    /// # Errors
    /// Returns an error if a valid field element could not be generated after 1000 calls to the
    /// PRNG.
    fn draw<E: FieldElement<BaseField = Self::BaseField>>(&mut self) -> Result<E, RandomCoinError>;

    /// Returns a vector of unique integers selected from the range [0, domain_size).
    ///
    /// # Errors
    /// Returns an error if the specified number of unique integers could not be generated
    /// after 1000 calls to the PRNG.
    ///
    /// # Panics
    /// Panics if:
    /// - `domain_size` is not a power of two.
    /// - `num_values` is greater than or equal to `domain_size`.
    fn draw_integers(
        &mut self,
        num_values: usize,
        domain_size: usize,
    ) -> Result<Vec<usize>, RandomCoinError>;
}


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

use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use examples::{fibonacci, Example};
use std::time::Duration;
use winterfell::{
    crypto::hashers::Blake3_256, math::fields::f128::BaseElement, FieldExtension, ProofOptions,
};

const SIZES: [usize; 3] = [16_384, 65_536, 262_144];

fn fibonacci(c: &mut Criterion) {
    let mut group = c.benchmark_group("fibonacci");
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(20));

    let options = ProofOptions::new(32, 8, 0, FieldExtension::None, 4, 255);

    for &size in SIZES.iter() {
        let fib =
            fibonacci::fib2::FibExample::<Blake3_256<BaseElement>>::new(size, options.clone());
        group.bench_function(BenchmarkId::from_parameter(size), |bench| {
            bench.iter(|| fib.prove());
        });
    }
    group.finish();
}

criterion_group!(fibonacci_group, fibonacci);
criterion_main!(fibonacci_group);


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

use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use examples::{rescue, Example};

use std::time::Duration;
use winterfell::{
    crypto::hashers::Blake3_256, math::fields::f128::BaseElement, FieldExtension, ProofOptions,
};

const SIZES: [usize; 2] = [256, 512];

fn rescue(c: &mut Criterion) {
    let mut group = c.benchmark_group("rescue");
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(25));

    let options = ProofOptions::new(32, 32, 0, FieldExtension::None, 4, 255);

    for &size in SIZES.iter() {
        let resc = rescue::RescueExample::<Blake3_256<BaseElement>>::new(size, options.clone());
        group.bench_function(BenchmarkId::from_parameter(size), |bench| {
            bench.iter(|| resc.prove());
        });
    }
    group.finish();
}

criterion_group!(rescue_group, rescue);
criterion_main!(rescue_group);