arcis-compiler 0.14.1

use crate::core::circuits::boolean::{
    boolean_value::Boolean,
    byte::Byte,
    utils::{addition_circuit, CircuitType},
};
use core::panic;
use std::ops::{Add, BitAnd, BitXor, Not, Shr};

/// u32 as a little-endian 4-Byte<B> array.
#[derive(Clone, Copy)]
pub(crate) struct U32<B: Boolean>([Byte<B>; 4]);

impl<B: Boolean> U32<B> {
    #[allow(dead_code)]
    pub fn from_le_bytes(bytes: [Byte<B>; 4]) -> Self {
        Self(bytes)
    }

    pub fn from_be_bytes(bytes: [Byte<B>; 4]) -> Self {
        let mut bytes = bytes;
        bytes.reverse();
        Self(bytes)
    }

    #[allow(dead_code)]
    pub fn to_le_bytes(self) -> [Byte<B>; 4] {
        self.0
    }

    pub fn to_be_bytes(self) -> [Byte<B>; 4] {
        let mut bytes = self.0;
        bytes.reverse();
        bytes
    }

    fn from_le_bits(bits: [B; 32]) -> Self {
        let mut res = [Byte::<B>::from(0u8); 4];
        for (i, byte) in bits
            .chunks(8)
            .map(|chunk| {
                let mut bits = [B::from(false); 8];
                bits.copy_from_slice(chunk);
                Byte::new(bits)
            })
            .enumerate()
        {
            res[i] = byte;
        }
        Self(res)
    }

    fn to_le_bits(self) -> [B; 32] {
        let mut res = [B::from(false); 32];
        for (i, byte) in self.0.iter().enumerate() {
            res[8 * i..8 * (i + 1)].copy_from_slice(&byte.get_bits());
        }
        res
    }

    pub fn rotr(&self, n: usize) -> Self {
        let n = n % 32;
        // lsb-to-msb representation of self
        let mut bits = self.to_le_bits().to_vec();
        bits.rotate_left(n);
        Self::from_le_bits(
            bits.try_into().unwrap_or_else(|v: Vec<B>| {
                panic!("Expected a Vec of length 32 (found {})", v.len())
            }),
        )
    }
}

impl<B: Boolean> BitXor for U32<B> {
    type Output = Self;

    fn bitxor(self, rhs: Self) -> Self::Output {
        let mut res = [Byte::<B>::from(0u8); 4];
        for (i, (lhs, rhs)) in self.0.into_iter().zip(rhs.0).enumerate() {
            res[i] = lhs ^ rhs;
        }
        Self(res)
    }
}

impl<B: Boolean> BitAnd for U32<B> {
    type Output = Self;

    fn bitand(self, rhs: Self) -> Self::Output {
        let mut res = [Byte::<B>::from(0u8); 4];
        for (i, (lhs, rhs)) in self.0.into_iter().zip(rhs.0).enumerate() {
            res[i] = lhs & rhs;
        }
        Self(res)
    }
}

impl<B: Boolean> Not for U32<B> {
    type Output = Self;

    fn not(self) -> Self::Output {
        let mut res = [Byte::<B>::from(0u8); 4];
        for (i, byte) in self.0.into_iter().enumerate() {
            res[i] = !byte;
        }
        Self(res)
    }
}

impl<B: Boolean> Add for U32<B> {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        // lsb-to-msb representations
        let [self_bits, rhs_bits] = [self, rhs].map(|value| value.to_le_bits());
        let mut sum = [B::from(false); 32];
        sum.copy_from_slice(&addition_circuit(
            self_bits.to_vec(),
            rhs_bits.to_vec(),
            B::from(false),
            CircuitType::default(),
        ));
        Self::from_le_bits(sum)
    }
}

impl<B: Boolean> Shr<usize> for U32<B> {
    type Output = Self;

    fn shr(self, rhs: usize) -> Self::Output {
        if rhs > 32 {
            panic!("shr by {rhs} positions not supported");
        } else if rhs == 32 {
            Self::from(0u32)
        } else {
            // lsb-to-msb representation of self[rhs..32]
            let mut bits = self.to_le_bits()[rhs..].to_vec();
            bits.resize(32, B::from(false));
            Self::from_le_bits(bits.try_into().unwrap_or_else(|v: Vec<B>| {
                panic!("Expected a Vec of length 32 (found {})", v.len())
            }))
        }
    }
}

impl<B: Boolean> From<u32> for U32<B> {
    fn from(value: u32) -> Self {
        let bytes = value.to_le_bytes();
        Self(bytes.map(Byte::<B>::from))
    }
}