arcis-compiler 0.9.7

use crate::{
    core::{
        actually_used_field::ActuallyUsedField,
        global_value::{field_array::FieldArray, value::FieldValue},
    },
    traits::{Pow, Random},
    utils::{matrix::Matrix, number::Number, used_field::UsedField},
};
use num_bigint::BigInt;
use num_traits::{ToPrimitive, Zero};
use sha3::{
    digest::{ExtendableOutput, Update, XofReader},
    Shake256,
};
use std::{
    fmt::Debug,
    iter::successors,
    ops::{Add, Mul, Sub},
};

pub trait RescueArg<F: UsedField>:
    Copy
    + Debug
    + Add<Self, Output = Self>
    + Sub<Self, Output = Self>
    + Mul<Self, Output = Self>
    + Mul<F, Output = Self>
    + Zero
    + Pow
    + Random
    + From<F>
{
}

impl<F: ActuallyUsedField> RescueArg<F> for F {}

impl<F: ActuallyUsedField> RescueArg<F> for FieldValue<F> {}

impl<const N: usize, F: ActuallyUsedField> RescueArg<F> for FieldArray<N, F> {}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum RescueMode {
    BlockCipher,
    HashFunction { capacity: usize },
}

// Security level for the block cipher.
const SECURITY_LEVEL_BLOCK_CIPHER: usize = 128;

// Security level for the hash function.
const SECURITY_LEVEL_HASH_FUNCTION: usize = 256;

/// see <https://tosc.iacr.org/index.php/ToSC/article/view/8695/8287> for everything
/// We used different MDS matrix: Cauchy-based matrix.
/// It's easier to compute and easier to prove it is MDS.

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct RescueDesc<F: UsedField, T: RescueArg<F>> {
    pub mode: RescueMode,
    alpha: Number,
    alpha_inverse: Number,
    n_rounds: usize,
    pub m: usize,
    mds_mat: Matrix<F>,
    mds_mat_inverse: Matrix<F>,
    round_keys: Vec<Matrix<T>>,
}

impl<F: UsedField, T: RescueArg<F>> RescueDesc<F, T> {
    pub fn new_cipher_desc(key: Matrix<T>) -> Self {
        if key.nrows == 1 || key.ncols > 1 {
            panic!(
                "key must be a column vector with at least 2 rows (found nrows: {}, ncols: {})",
                key.nrows, key.ncols
            );
        }
        let m = key.nrows;
        let alpha = F::get_alpha();
        let alpha_inverse = F::get_alpha_inverse();
        let n_rounds = Self::get_n_rounds(RescueMode::BlockCipher, &alpha, m);
        let (mds_mat, mds_mat_inverse) = F::mds_matrix_and_inverse(m);
        // generate the round constants using SHAKE256 hash
        let round_constants = Self::sample_constants(RescueMode::BlockCipher, n_rounds, m);

        // do the key schedule
        let round_keys = rescue_permutation(
            RescueMode::BlockCipher,
            &alpha,
            &alpha_inverse,
            &mds_mat,
            &round_constants
                .into_iter()
                .map(|c| c.convert())
                .collect::<Vec<Matrix<T>>>(),
            &key,
        );

        RescueDesc {
            mode: RescueMode::BlockCipher,
            alpha,
            alpha_inverse,
            n_rounds,
            m,
            mds_mat,
            mds_mat_inverse,
            round_keys,
        }
    }

    pub fn new_hash_desc(m: usize, capacity: usize) -> Self {
        let alpha = F::get_alpha();
        let alpha_inverse = F::get_alpha_inverse();
        let n_rounds = Self::get_n_rounds(RescueMode::HashFunction { capacity }, &alpha, m);
        let (mds_mat, mds_mat_inverse) = F::mds_matrix_and_inverse(m);
        // generate the round constants using SHAKE256 hash
        let round_constants =
            Self::sample_constants(RescueMode::HashFunction { capacity }, n_rounds, m);

        RescueDesc {
            mode: RescueMode::HashFunction { capacity },
            alpha,
            alpha_inverse,
            n_rounds,
            m,
            mds_mat,
            mds_mat_inverse,
            round_keys: round_constants
                .into_iter()
                .map(|c| c.convert())
                .collect::<Vec<Matrix<T>>>(),
        }
    }

    fn get_n_rounds(mode: RescueMode, alpha: &Number, m: usize) -> usize {
        match mode {
            RescueMode::BlockCipher => {
                // see https://tosc.iacr.org/index.php/ToSC/article/view/8695/8287
                let l_0 = (SECURITY_LEVEL_BLOCK_CIPHER as f64 * 2.0
                    / ((m + 1) as f64 * (F::modulus().log2() - (alpha - 1).log2())))
                .ceil() as usize;
                let l_1 = if *alpha == 3 {
                    ((SECURITY_LEVEL_BLOCK_CIPHER + 2) as f64 / (4 * m) as f64).ceil() as usize
                } else {
                    ((SECURITY_LEVEL_BLOCK_CIPHER + 3) as f64 / (m as f64 * 5.5)).ceil() as usize
                };
                2 * ([l_0, l_1, 5].into_iter().max().unwrap())
            }
            RescueMode::HashFunction { capacity } => {
                // get number of rounds for Groebner basis attack
                let rate = m - capacity;
                fn dcon(n: usize, alpha: &Number, m: usize) -> usize {
                    (0.5 * ((alpha.to_usize().unwrap() - 1) * m * (n - 1)) as f64 + 2.0).floor()
                        as usize
                }
                fn v(n: usize, rate: usize, m: usize) -> usize {
                    m * (n - 1) + rate
                }
                fn binomial(n: usize, k: usize) -> Number {
                    fn factorial(m: Number) -> Number {
                        if m == 0 || m == 1 {
                            Number::from(1)
                        } else {
                            m.clone() * factorial(m - 1)
                        }
                    }
                    factorial(Number::from(n))
                        / (factorial(Number::from(n - k)) * factorial(Number::from(k)))
                }

                let target = Number::power_of_two(SECURITY_LEVEL_HASH_FUNCTION);
                let mut l1 = 1;
                let mut tmp = binomial(v(l1, rate, m) + dcon(l1, alpha, m), v(l1, rate, m));
                while tmp.clone() * tmp <= target && l1 <= 23 {
                    l1 += 1;
                    tmp = binomial(v(l1, rate, m) + dcon(l1, alpha, m), v(l1, rate, m));
                }

                // set a minimum value for sanity and add 50%
                (1.5 * [5, l1].into_iter().max().unwrap() as f64).ceil() as usize
            }
        }
    }

    fn sample_constants(mode: RescueMode, n_rounds: usize, m: usize) -> Vec<Matrix<F>> {
        // setup randomness
        let mut hasher = Shake256::default();
        // buffer to create `FieldElements` from bytes (via `Number`)
        // we add 16 bytes to get a distribution statistically close to uniform
        let buffer_len = F::NUM_BITS.div_ceil(8) as usize + 16;
        match mode {
            RescueMode::BlockCipher => {
                hasher.update(b"encrypt everything, compute anything");
                let mut reader = hasher.finalize_xof();

                let mut f_iter = (0..m * m + 2 * m).map(|_| {
                    // create field element from the shake hash
                    let randomness = reader.read_boxed(buffer_len);
                    // we set the sign to plus to essentially read unsigned BigInts (BigUInts),
                    // matching the noble curves TypeScript implementation used
                    // in the client.
                    let b = BigInt::from_bytes_le(num_bigint::Sign::Plus, &randomness);
                    // we need not check whether the obtained field element f is in any subgroup,
                    // because we use only prime fields (i.e. there are no subgroups)
                    F::from(Number::from(b))
                });

                // create matrix and vectors
                let mut round_constant_mat =
                    Matrix::new_from_iter((m, m), (&mut f_iter).take(m * m));
                let initial_round_constant = Matrix::new_from_iter((m, 1), (&mut f_iter).take(m));
                let round_constant_affine_term =
                    Matrix::new_from_iter((m, 1), (&mut f_iter).take(m));

                // check for inversability
                while round_constant_mat.det() == F::ZERO {
                    //resample the matrix
                    let data = vec![F::ZERO; m * m].into_iter().map(|_| {
                        let randomness = reader.read_boxed(buffer_len);
                        let b = BigInt::from_bytes_le(num_bigint::Sign::Plus, &randomness);
                        F::from(Number::from(b))
                    });
                    round_constant_mat = Matrix::new_from_iter((m, m), data);
                }

                let mut iter = 0..2 * n_rounds;
                successors(Some(initial_round_constant), |c| {
                    iter.next().map(|_| {
                        round_constant_mat.clone().mat_mul(c) + round_constant_affine_term.clone()
                    })
                })
                .collect::<Vec<Matrix<F>>>()
            }
            RescueMode::HashFunction { capacity } => {
                let seed = format!(
                    "Rescue-XLIX({},{},{},{})",
                    F::modulus(),
                    m,
                    capacity,
                    SECURITY_LEVEL_HASH_FUNCTION
                );
                hasher.update(seed.as_bytes());
                let mut reader = hasher.finalize_xof();

                let mut round_constants = (0..2 * m * n_rounds)
                    .map(|_| {
                        // create field element from the shake hash
                        let randomness = reader.read_boxed(buffer_len);
                        // we set the sign to plus to essentially read unsigned BigInts (BigUInts),
                        // matching the noble curves TypeScript implementation used
                        // in the client.
                        let b = BigInt::from_bytes_le(num_bigint::Sign::Plus, &randomness);
                        // we need not check whether the obtained field element f is in any
                        // subgroup, because we use only prime fields (i.e.
                        // there are no subgroups)
                        F::from(Number::from(b))
                    })
                    .collect::<Vec<F>>()
                    .chunks(m)
                    .map(|c| Matrix::new_from_iter((m, 1), c.iter().copied()))
                    .collect::<Vec<Matrix<F>>>();
                // Self::permute requires an odd number of round keys
                // prepending a 0 matrix makes it equivalent to Algorithm 3 from https://eprint.iacr.org/2020/1143.pdf
                round_constants.insert(0, Matrix::new((m, 1), F::ZERO));
                round_constants
            }
        }
    }

    pub fn permute(&self, state: &Matrix<T>) -> Matrix<T> {
        rescue_permutation(
            self.mode,
            &self.alpha,
            &self.alpha_inverse,
            &self.mds_mat,
            &self.round_keys,
            state,
        )
        .last()
        .unwrap()
        .clone()
    }

    pub fn permute_inverse(&self, state: &Matrix<T>) -> Matrix<T> {
        rescue_permutation_inverse(
            self.mode,
            &self.alpha,
            &self.alpha_inverse,
            &self.mds_mat_inverse,
            &self.round_keys,
            state,
        )
        .last()
        .unwrap()
        .clone()
    }
}

fn exponent_for_even(mode: RescueMode, alpha: Number, alpha_inverse: Number) -> Number {
    match mode {
        RescueMode::BlockCipher => alpha_inverse,
        RescueMode::HashFunction { capacity: _ } => alpha,
    }
}

fn exponent_for_odd(mode: RescueMode, alpha: Number, alpha_inverse: Number) -> Number {
    match mode {
        RescueMode::BlockCipher => alpha,
        RescueMode::HashFunction { capacity: _ } => alpha_inverse,
    }
}

fn rescue_permutation<T: RescueArg<F>, F: UsedField>(
    mode: RescueMode,
    alpha: &Number,
    alpha_inverse: &Number,
    mds_mat: &Matrix<F>,
    subkeys: &[Matrix<T>],
    state: &Matrix<T>,
) -> Vec<Matrix<T>> {
    let exponent_even = exponent_for_even(mode, alpha.clone(), alpha_inverse.clone());
    let exponent_odd = exponent_for_odd(mode, alpha.clone(), alpha_inverse.clone());
    let initial_key = &subkeys[0];
    let mut iter = subkeys[1..].iter().enumerate();
    successors(Some(state.clone() + initial_key.clone()), |s| {
        iter.next().map(|(r, key)| {
            let mut s = s.clone();
            if r % 2 == 0 {
                // we can expect x to be non-zero
                s.map_mut(|x| x.pow(&exponent_even, true));
            } else {
                // we can expect x to be non-zero
                s.map_mut(|x| x.pow(&exponent_odd, true));
            }
            s = mds_mat.mat_mul(&s);
            s += key;
            s
        })
    })
    .collect::<Vec<Matrix<T>>>()
}

fn rescue_permutation_inverse<T: RescueArg<F>, F: UsedField>(
    mode: RescueMode,
    alpha: &Number,
    alpha_inverse: &Number,
    mds_mat_inverse: &Matrix<F>,
    subkeys: &[Matrix<T>],
    state: &Matrix<T>,
) -> Vec<Matrix<T>> {
    let exponent_even = exponent_for_even(mode, alpha.clone(), alpha_inverse.clone());
    let exponent_odd = exponent_for_odd(mode, alpha.clone(), alpha_inverse.clone());
    let initial_key = &subkeys[0];
    let mut states = subkeys[1..]
        .iter()
        .rev()
        .enumerate()
        .scan(state.clone(), |s, (r, key)| {
            *s -= key;
            *s = mds_mat_inverse.mat_mul(s);
            if r % 2 == 0 {
                // we can expect x to be non-zero
                s.map_mut(|x| x.pow(&exponent_even, true));
            } else {
                // we can expect x to be non-zero
                s.map_mut(|x| x.pow(&exponent_odd, true));
            }
            Some(s.clone())
        })
        .collect::<Vec<Matrix<T>>>();
    states.push(states.last().unwrap().clone() - initial_key.clone());
    states
}
#[cfg(test)]
mod tests {
    use super::*;
    use crate::utils::field::BaseField;
    use ff::Field;
    use rand::Rng;

    fn test_rescue_desc<R: Rng + ?Sized>(rng: &mut R, rescue: RescueDesc<BaseField, BaseField>) {
        let alpha_prod = (&rescue.alpha * &rescue.alpha_inverse) % (BaseField::modulus() - 1);
        assert_eq!(alpha_prod, Number::from(1));
        fn test_is_identity(mat_prod: Matrix<BaseField>) {
            for i in 0..mat_prod.nrows {
                for j in 0..mat_prod.ncols {
                    let expected = if i == j {
                        BaseField::ONE
                    } else {
                        BaseField::ZERO
                    };
                    assert_eq!(*mat_prod.get((i, j)).unwrap(), expected);
                }
            }
        }
        let mat_prod = rescue.mds_mat.mat_mul(&rescue.mds_mat_inverse);
        test_is_identity(mat_prod);
        let mat_prod = rescue.mds_mat_inverse.mat_mul(&rescue.mds_mat);
        test_is_identity(mat_prod);
        for _ in 0..2 {
            let state = Matrix::from(gen_random_fp(rng, rescue.m));
            let permuted = rescue.permute(&state);
            let unpermuted = rescue.permute_inverse(&permuted);
            assert_eq!(unpermuted, state);
        }
    }
    fn gen_random_fp<R: Rng + ?Sized>(rng: &mut R, size: usize) -> Vec<BaseField> {
        (0..size)
            .map(|_| <BaseField as ff::Field>::random(&mut *rng))
            .collect()
    }
    #[test]
    fn rescue_desc() {
        let rng = &mut crate::utils::test_rng::get();

        let mut m = 2;
        while rng.gen_bool(0.5) {
            m += 1;
        }
        let rescue_cipher = RescueDesc::new_cipher_desc(Matrix::from(gen_random_fp(rng, m)));
        test_rescue_desc(rng, rescue_cipher);

        let mut capacity = 1;
        while rng.gen_bool(0.5) {
            capacity += 1;
        }
        capacity = capacity.min(m - 1);
        let rescue_hash = RescueDesc::<BaseField, BaseField>::new_hash_desc(m, capacity);
        test_rescue_desc(rng, rescue_hash);
    }
}