arcis-compiler 0.9.7

use crate::{
    core::{
        actually_used_field::ActuallyUsedField,
        circuits::{
            arithmetic::PowCircuit,
            boolean::{boolean_array::BooleanArray, boolean_value::BooleanValue, byte::Byte},
            traits::arithmetic_circuit::ArithmeticCircuit,
        },
        global_value::value::FieldValue,
    },
    traits::{FromF25519, Invert, Pow, Random, Reveal, ToLeBytes, WithBooleanBounds},
    utils::{
        field::{BaseField, ScalarField},
        number::Number,
    },
};
use num_traits::Zero;
use std::ops::{Add, Index, Mul, Neg, Sub};

#[derive(Debug, Clone, Copy)]
pub struct FieldArray<const N: usize, F: ActuallyUsedField>([FieldValue<F>; N]);

impl<const N: usize, F: ActuallyUsedField> IntoIterator for FieldArray<N, F> {
    type Item = FieldValue<F>;
    type IntoIter = <[FieldValue<F>; N] as IntoIterator>::IntoIter;

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

impl<const N: usize, F: ActuallyUsedField> Index<usize> for FieldArray<N, F> {
    type Output = FieldValue<F>;

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

impl<const N: usize, F: ActuallyUsedField> Add for FieldArray<N, F> {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self(
            self.into_iter()
                .zip(rhs)
                .map(|(val_lhs, val_rhs)| val_lhs + val_rhs)
                .collect::<Vec<FieldValue<F>>>()
                .try_into()
                .unwrap_or_else(|v: Vec<FieldValue<F>>| {
                    panic!("Expected a Vec of length {} (found {})", N, v.len())
                }),
        )
    }
}

impl<const N: usize, F: ActuallyUsedField> Sub for FieldArray<N, F> {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        Self(
            self.into_iter()
                .zip(rhs)
                .map(|(val_lhs, val_rhs)| val_lhs - val_rhs)
                .collect::<Vec<FieldValue<F>>>()
                .try_into()
                .unwrap_or_else(|v: Vec<FieldValue<F>>| {
                    panic!("Expected a Vec of length {} (found {})", N, v.len())
                }),
        )
    }
}

impl<const N: usize, F: ActuallyUsedField> Mul for FieldArray<N, F> {
    type Output = Self;

    fn mul(self, rhs: Self) -> Self::Output {
        Self(
            self.into_iter()
                .zip(rhs)
                .map(|(val_lhs, val_rhs)| val_lhs * val_rhs)
                .collect::<Vec<FieldValue<F>>>()
                .try_into()
                .unwrap_or_else(|v: Vec<FieldValue<F>>| {
                    panic!("Expected a Vec of length {} (found {})", N, v.len())
                }),
        )
    }
}

impl<const N: usize, F: ActuallyUsedField> Neg for FieldArray<N, F> {
    type Output = Self;

    fn neg(self) -> Self::Output {
        Self(self.0.map(|val| -val))
    }
}

impl<const N: usize, F: ActuallyUsedField> Mul<F> for FieldArray<N, F> {
    type Output = Self;

    fn mul(self, rhs: F) -> Self::Output {
        Self(
            self.into_iter()
                .map(|val| val * rhs)
                .collect::<Vec<FieldValue<F>>>()
                .try_into()
                .unwrap_or_else(|v: Vec<FieldValue<F>>| {
                    panic!("Expected a Vec of length {} (found {})", N, v.len())
                }),
        )
    }
}

impl<const N: usize, F: ActuallyUsedField> Zero for FieldArray<N, F> {
    fn zero() -> Self {
        Self::from(0)
    }

    fn is_zero(&self) -> bool {
        self.into_iter()
            .map(|val| val.is_zero())
            .reduce(|a, b| a & b)
            .unwrap()
    }
}

impl<const N: usize, F: ActuallyUsedField> Pow for FieldArray<N, F> {
    fn pow(self, e: &Number, is_expected_non_zero: bool) -> Self {
        let e = e % (F::modulus() - 1);
        let pow_circuit = PowCircuit::new(e, is_expected_non_zero);
        // TODO: implement PowCircuit for FieldArray
        Self(self.0.map(|val| pow_circuit.run(vec![val])[0]))
    }
}

impl<const N: usize, F: ActuallyUsedField> Invert for FieldArray<N, F> {
    fn invert(self, is_expected_non_zero: bool) -> Self {
        Self(self.0.map(|val| val.invert(is_expected_non_zero)))
    }
}

impl<const N: usize, F: ActuallyUsedField> Reveal for FieldArray<N, F> {
    fn reveal(self) -> Self {
        Self(self.0.map(|val| val.reveal()))
    }
}

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

impl<const N: usize, F: ActuallyUsedField> WithBooleanBounds for FieldArray<N, F> {
    fn with_boolean_bounds(&self) -> Self {
        todo!()
    }
}

impl<const N: usize, F: ActuallyUsedField> From<[FieldValue<F>; N]> for FieldArray<N, F> {
    fn from(value: [FieldValue<F>; N]) -> Self {
        Self(value)
    }
}

impl<const N: usize, F: ActuallyUsedField> From<FieldArray<N, F>> for [FieldValue<F>; N] {
    fn from(value: FieldArray<N, F>) -> Self {
        value.0
    }
}

impl<const N: usize, F: ActuallyUsedField> From<i32> for FieldArray<N, F> {
    fn from(value: i32) -> Self {
        Self([FieldValue::<F>::from(value); N])
    }
}

impl<const N: usize, F: ActuallyUsedField> From<Number> for FieldArray<N, F> {
    fn from(value: Number) -> Self {
        Self([FieldValue::<F>::from(value); N])
    }
}

impl<const N: usize, F: ActuallyUsedField> From<F> for FieldArray<N, F> {
    fn from(value: F) -> Self {
        Self([FieldValue::from(value); N])
    }
}

impl<const N: usize, F: ActuallyUsedField> From<FieldValue<F>> for FieldArray<N, F> {
    fn from(value: FieldValue<F>) -> Self {
        Self([value; N])
    }
}

impl<const N: usize, F: ActuallyUsedField> From<BooleanArray<N>> for FieldArray<N, F> {
    fn from(value: BooleanArray<N>) -> Self {
        Self(
            value
                .into_iter()
                .map(FieldValue::from)
                .collect::<Vec<FieldValue<F>>>()
                .try_into()
                .unwrap_or_else(|v: Vec<FieldValue<F>>| {
                    panic!("Expected a Vec of length {} (found {})", N, v.len())
                }),
        )
    }
}

impl<const N: usize, F: ActuallyUsedField> ToLeBytes for FieldArray<N, F> {
    type BooleanOutput = BooleanArray<N>;

    fn to_le_bytes(self) -> [Byte<Self::BooleanOutput>; 32] {
        let bytes_vec = self
            .0
            .into_iter()
            .map(|val| val.to_le_bytes())
            .collect::<Vec<[Byte<BooleanValue>; 32]>>();

        (0..32)
            .map(|i| {
                Byte::new(
                    (0..8)
                        .map(|j| {
                            BooleanArray::from(
                                TryInto::<[BooleanValue; N]>::try_into(
                                    bytes_vec
                                        .iter()
                                        .map(|bytes| bytes[i].get_bits()[j])
                                        .collect::<Vec<BooleanValue>>(),
                                )
                                .unwrap_or_else(
                                    |v: Vec<BooleanValue>| {
                                        panic!("Expected a Vec of length {} (found {})", N, v.len())
                                    },
                                ),
                            )
                        })
                        .collect::<Vec<BooleanArray<N>>>()
                        .try_into()
                        .unwrap_or_else(|v: Vec<BooleanArray<N>>| {
                            panic!("Expected a Vec of length 8 (found {})", v.len())
                        }),
                )
            })
            .collect::<Vec<Byte<Self::BooleanOutput>>>()
            .try_into()
            .unwrap_or_else(|v: Vec<Byte<Self::BooleanOutput>>| {
                panic!("Expected a Vec of length 32 (found {})", v.len())
            })
    }
}

impl<const N: usize> FromF25519<FieldArray<N, BaseField>> for FieldArray<N, BaseField> {
    fn from_F25519(value: FieldArray<N, BaseField>) -> Vec<FieldArray<N, BaseField>> {
        vec![value]
    }
}

impl<const N: usize> FromF25519<FieldArray<N, BaseField>> for FieldArray<N, ScalarField> {
    fn from_F25519(_value: FieldArray<N, BaseField>) -> Vec<FieldArray<N, ScalarField>> {
        // let bytes = value.to_le_bytes();
        // // we chunk bytes by ScalarField::n_bytes - 1 and convert to FieldValue<ScalarField>
        // let chunk_size = ScalarField::NUM_BITS.div_ceil(8) as usize - 1;
        // bytes
        //     .chunks(chunk_size)
        //     .map(|chunk| {
        //         FieldArray::<N, ScalarField>::from_le_bits(
        //             chunk
        //                 .iter()
        //                 .flat_map(|byte| byte.to_vec())
        //                 .collect::<Vec<BooleanArray<N>>>(),
        //             false,
        //         )
        //     })
        //     .collect::<Vec<FieldArray<N, ScalarField>>>()
        todo!();
    }
}