arcis-compiler 0.14.1

use crate::{
    core::{
        bounds::FieldBounds,
        circuits::{
            boolean::{
                boolean_value::{Boolean, BooleanValue},
                byte::Byte,
                u32::U32,
            },
            traits::arithmetic_circuit::ArithmeticCircuit,
        },
        expressions::expr::EvalFailure,
        global_value::value::FieldValue,
    },
    utils::field::BaseField,
};
use sha2::Digest;

#[derive(Debug, Clone)]
pub struct Sha256;

impl Sha256 {
    pub const DIGEST_BYTES: usize = 32;

    pub const BLOCK_BYTES: usize = 64;

    const K: [u32; 64] = [
        1116352408, 1899447441, 3049323471, 3921009573, 961987163, 1508970993, 2453635748,
        2870763221, 3624381080, 310598401, 607225278, 1426881987, 1925078388, 2162078206,
        2614888103, 3248222580, 3835390401, 4022224774, 264347078, 604807628, 770255983,
        1249150122, 1555081692, 1996064986, 2554220882, 2821834349, 2952996808, 3210313671,
        3336571891, 3584528711, 113926993, 338241895, 666307205, 773529912, 1294757372, 1396182291,
        1695183700, 1986661051, 2177026350, 2456956037, 2730485921, 2820302411, 3259730800,
        3345764771, 3516065817, 3600352804, 4094571909, 275423344, 430227734, 506948616, 659060556,
        883997877, 958139571, 1322822218, 1537002063, 1747873779, 1955562222, 2024104815,
        2227730452, 2361852424, 2428436474, 2756734187, 3204031479, 3329325298,
    ];

    const H: [u32; 8] = [
        1779033703, 3144134277, 1013904242, 2773480762, 1359893119, 2600822924, 528734635,
        1541459225,
    ];

    pub fn new() -> Self {
        Self
    }

    pub fn digest<B: Boolean>(&self, message: Vec<Byte<B>>) -> [Byte<B>; Self::DIGEST_BYTES] {
        let len = message.len();
        let bitlen_bytes = (8 * len as u64).to_be_bytes();
        let l = len % 64;
        let padding = if l < 56 {
            64 - (l + 1 + 8)
        } else {
            128 - (l + 1 + 8)
        };
        let mut message_padded = message;
        message_padded.push(Byte::<B>::from(128u8));
        message_padded.extend(vec![Byte::<B>::from(0u8); padding]);
        message_padded.extend(bitlen_bytes.map(Byte::<B>::from));

        let mut self_h = Self::H.map(U32::<B>::from);
        let self_k = Self::K.map(U32::<B>::from);

        for chunk in message_padded.chunks(64) {
            let mut w = [U32::<B>::from(0u32); 64];
            let chunk_u32 = chunk
                .chunks(4)
                .map(|c| {
                    let mut bytes = [Byte::<B>::from(0u8); 4];
                    for (i, byte) in c.iter().enumerate() {
                        bytes[i] = *byte;
                    }
                    U32::from_be_bytes(bytes)
                })
                .collect::<Vec<U32<B>>>();
            w[..16].copy_from_slice(&chunk_u32);
            for i in 16..64 {
                let s0 = w[i - 15].rotr(7) ^ w[i - 15].rotr(18) ^ (w[i - 15] >> 3);
                let s1 = w[i - 2].rotr(17) ^ w[i - 2].rotr(19) ^ (w[i - 2] >> 10);
                w[i] = (w[i - 16] + s0) + (w[i - 7] + s1);
            }
            let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h] = self_h;

            for i in 0..64 {
                let s0 = a.rotr(2) ^ a.rotr(13) ^ a.rotr(22);
                let maj = (a & b) ^ (a & c) ^ (b & c);
                let t2 = s0 + maj;
                let s1 = e.rotr(6) ^ e.rotr(11) ^ e.rotr(25);
                let ch = (e & f) ^ ((!e) & g);
                // We have t1 = h + s1 + ch + self_k[i] + w[i], t2 = s0 + maj,
                // e = d + t1 and a = t1 + t2.
                // Without setting parantheses, t1 will be computed in 4 consecutive
                // adders, while e and a will be computed in 5 and 6 consecutive
                // adders respectively. We slightly re-organize the terms, set parantheses
                // to enforce parallel adders, and reduce the number of consecutive
                // adders to 3.
                // two consecutive adders
                let t1_tmp = (h + s1) + (ch + self_k[i]);
                // one adder
                let d_tmp = d + w[i];
                // two consecutive adders
                let t2_tmp = w[i] + t2;

                h = g;
                g = f;
                f = e;
                e = d_tmp + t1_tmp;
                d = c;
                c = b;
                b = a;
                a = t1_tmp + t2_tmp;
            }

            for (lhs, rhs) in self_h.iter_mut().zip([a, b, c, d, e, f, g, h]) {
                *lhs = *lhs + rhs;
            }
        }

        let mut digest = [Byte::<B>::from(0u8); Self::DIGEST_BYTES];
        for (i, value) in self_h.into_iter().enumerate() {
            digest[4 * i..4 * (i + 1)].copy_from_slice(&value.to_be_bytes());
        }
        digest
    }
}

impl Default for Sha256 {
    fn default() -> Self {
        Self::new()
    }
}

impl ArithmeticCircuit<BaseField> for Sha256 {
    fn eval(&self, x: Vec<BaseField>) -> Result<Vec<BaseField>, EvalFailure> {
        // all inputs are expected to be bytes
        x.iter()
            .for_each(|byte| assert!(*byte <= BaseField::from(255)));
        let message = x
            .into_iter()
            .map(|val| val.to_le_bytes()[0])
            .collect::<Vec<u8>>();

        let mut hasher = sha2::Sha256::new();
        hasher.update(message);
        let digest = hasher.finalize();
        Ok(digest
            .iter()
            .map(|byte| BaseField::from(*byte as u64))
            .collect::<Vec<BaseField>>())
    }

    fn bounds(&self, _bounds: Vec<FieldBounds<BaseField>>) -> Vec<FieldBounds<BaseField>> {
        vec![FieldBounds::new(BaseField::from(0), BaseField::from(255)); 32]
    }

    fn run(&self, vals: Vec<FieldValue<BaseField>>) -> Vec<FieldValue<BaseField>> {
        let message = vals
            .into_iter()
            .map(Byte::from)
            .collect::<Vec<Byte<BooleanValue>>>();

        let hasher = Sha256::new();
        hasher
            .digest(message)
            .into_iter()
            .map(FieldValue::<BaseField>::from)
            .collect::<Vec<FieldValue<BaseField>>>()
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::circuits::traits::arithmetic_circuit::tests::TestedArithmeticCircuit;
    use rand::Rng;

    impl TestedArithmeticCircuit<BaseField> for Sha256 {
        fn gen_desc<R: Rng + ?Sized>(_rng: &mut R) -> Self {
            Self
        }

        fn gen_n_inputs<R: Rng + ?Sized>(&self, rng: &mut R) -> usize {
            (rng.next_u32() % 128).try_into().unwrap()
        }

        fn gen_input_bounds<R: Rng + ?Sized>(_rng: &mut R) -> FieldBounds<BaseField> {
            FieldBounds::new(BaseField::from(0), BaseField::from(255))
        }
    }

    #[test]
    fn tested_sha256() {
        Sha256::test(1, 1)
    }
}