arcis-compiler 0.9.6

use crate::{
    core::{
        actually_used_field::ActuallyUsedField,
        circuits::boolean::boolean_value::{Boolean, BooleanValue},
        global_value::{field_array::FieldArray, value::FieldValue},
    },
    traits::{RandomBit, Reveal, Select},
};
use core::panic;
use std::ops::{BitAnd, BitAndAssign, BitXor, BitXorAssign, Index, Not};

#[derive(Debug, Clone, Copy)]
pub struct BooleanArray<const N: usize>([BooleanValue; N]);

impl<const N: usize> IntoIterator for BooleanArray<N> {
    type Item = BooleanValue;
    type IntoIter = <[BooleanValue; N] as IntoIterator>::IntoIter;

    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl<const N: usize> Index<usize> for BooleanArray<N> {
    type Output = BooleanValue;

    fn index(&self, index: usize) -> &Self::Output {
        &self.0[index]
    }
}

impl<const N: usize> BitAnd for BooleanArray<N> {
    type Output = BooleanArray<N>;

    fn bitand(self, rhs: Self) -> Self::Output {
        Self(
            self.into_iter()
                .zip(rhs)
                .map(|(bit_lhs, bit_rhs)| bit_lhs & bit_rhs)
                .collect::<Vec<BooleanValue>>()
                .try_into()
                .unwrap_or_else(|v: Vec<BooleanValue>| {
                    panic!("Expected a Vec of length {} (found {})", N, v.len())
                }),
        )
    }
}

impl<const N: usize> BitAndAssign for BooleanArray<N> {
    fn bitand_assign(&mut self, rhs: Self) {
        *self = *self & rhs;
    }
}

impl<const N: usize> BitXor for BooleanArray<N> {
    type Output = BooleanArray<N>;

    fn bitxor(self, rhs: Self) -> Self::Output {
        Self(
            self.into_iter()
                .zip(rhs)
                .map(|(bit_lhs, bit_rhs)| bit_lhs ^ bit_rhs)
                .collect::<Vec<BooleanValue>>()
                .try_into()
                .unwrap_or_else(|v: Vec<BooleanValue>| {
                    panic!("Expected a Vec of length {} (found {})", N, v.len())
                }),
        )
    }
}

impl<const N: usize> BitXorAssign for BooleanArray<N> {
    fn bitxor_assign(&mut self, rhs: Self) {
        *self = *self ^ rhs;
    }
}

impl<const N: usize> Not for BooleanArray<N> {
    type Output = BooleanArray<N>;

    fn not(self) -> Self::Output {
        Self(self.0.map(|bit| !bit))
    }
}

impl<const N: usize> RandomBit for BooleanArray<N> {
    fn random() -> Self {
        Self(
            (0..N)
                .map(|_| BooleanValue::random())
                .collect::<Vec<BooleanValue>>()
                .try_into()
                .unwrap_or_else(|v: Vec<BooleanValue>| {
                    panic!("Expected a Vec of length {} (found {})", N, v.len())
                }),
        )
    }
}

impl<const N: usize> Reveal for BooleanArray<N> {
    fn reveal(self) -> Self {
        Self(self.0.map(|bit| bit.reveal()))
    }
}

impl<const N: usize> Select<BooleanArray<N>, BooleanArray<N>, BooleanArray<N>> for BooleanArray<N> {
    fn select(self, x: BooleanArray<N>, y: BooleanArray<N>) -> BooleanArray<N> {
        y ^ self & (x ^ y)
    }
}

impl<const N: usize> Select<Vec<BooleanArray<N>>, Vec<BooleanArray<N>>, Vec<BooleanArray<N>>>
    for BooleanArray<N>
{
    fn select(self, x: Vec<BooleanArray<N>>, y: Vec<BooleanArray<N>>) -> Vec<BooleanArray<N>> {
        assert_eq!(x.len(), y.len());
        x.into_iter()
            .zip(y)
            .map(|(bit_x, bit_y)| self.select(bit_x, bit_y))
            .collect::<Vec<BooleanArray<N>>>()
    }
}

impl<const N: usize, F: ActuallyUsedField>
    Select<FieldArray<N, F>, FieldArray<N, F>, FieldArray<N, F>> for BooleanArray<N>
{
    fn select(self, x: FieldArray<N, F>, y: FieldArray<N, F>) -> FieldArray<N, F> {
        FieldArray::from(
            TryInto::<[FieldValue<F>; N]>::try_into(
                self.into_iter()
                    .zip(x)
                    .zip(y)
                    .map(|((cond, val_x), val_y)| cond.select(val_x, val_y))
                    .collect::<Vec<FieldValue<F>>>(),
            )
            .unwrap_or_else(|v: Vec<FieldValue<F>>| {
                panic!("Expected a Vec of length {} (found {})", N, v.len())
            }),
        )
    }
}

impl<const N: usize> From<[BooleanValue; N]> for BooleanArray<N> {
    fn from(value: [BooleanValue; N]) -> Self {
        Self(value)
    }
}

impl<const N: usize> From<BooleanArray<N>> for [BooleanValue; N] {
    fn from(value: BooleanArray<N>) -> Self {
        value.0
    }
}

impl<const N: usize> From<bool> for BooleanArray<N> {
    fn from(value: bool) -> Self {
        Self([BooleanValue::from(value); N])
    }
}

impl<const N: usize> From<BooleanValue> for BooleanArray<N> {
    fn from(value: BooleanValue) -> Self {
        Self([value; N])
    }
}

impl<const N: usize> Boolean for BooleanArray<N> {}