arcis-compiler 0.9.7

use crate::{
    core::{
        actually_used_field::ActuallyUsedField,
        bounds::{FieldBounds, IsBounds},
        circuits::{
            arithmetic::PowCircuit,
            boolean::{boolean_value::BooleanValue, byte::Byte, utils::sign_bit},
            traits::arithmetic_circuit::ArithmeticCircuit,
        },
        expressions::{
            conversion_expr::ConversionExpr::{BitNumToBit, BitToBitNum, ScalarFromPlaintextBit},
            expr::Expr,
            field_expr::{FieldExpr, RandomValId},
            other_expr::OtherExpr,
        },
        global_value::{
            global_expr_store::{are_unbuilt_circuits_allowed, with_global_expr_store_as_local},
            value::FieldValue,
        },
        ir_builder::{ExprStore, IRBuilder},
    },
    traits::{
        Equal,
        FromF25519,
        FromLeBits,
        GetBit,
        GreaterEqual,
        GreaterThan,
        Invert,
        MxeRescueKey,
        MxeX25519PrivateKey,
        Pow,
        Random,
        RandomBit,
        Reveal,
        Select,
        Selectable,
        ToLeBytes,
        WithBooleanBounds,
    },
    utils::{
        crypto::key::RESCUE_KEY_COUNT,
        field::{BaseField, ScalarField},
        number::Number,
    },
    MxeFieldInput,
    MxeInput,
    MxeScalarInput,
};
use ff::PrimeField;
use num_traits::Zero;
use std::{
    iter::Sum,
    ops::{
        Add,
        AddAssign,
        Div,
        DivAssign,
        Mul,
        MulAssign,
        Neg,
        Not,
        Rem,
        RemAssign,
        Shr,
        Sub,
        SubAssign,
    },
};

impl<F: ActuallyUsedField> FieldValue<F> {
    /// The `Expr<usize>` which generates the value of `self`.
    pub fn expr(&self) -> Expr<usize> {
        with_global_expr_store_as_local(|expr_store| expr_store.get_expr(self.get_id()).clone())
    }

    /// Used when we know the value of `self` satisfy these bounds.
    /// The compiler will intersect these bounds with the ones it computed.
    /// Note: creates a new `FieldValue`, does not change `self`'s bounds.
    pub fn with_bounds(&self, bounds: impl Into<FieldBounds<F>>) -> FieldValue<F> {
        Self::new(FieldExpr::Bounds(*self, bounds.into()))
    }

    /// The computed bounds of this value.
    pub fn bounds(&self) -> FieldBounds<F> {
        F::bounds_to_field_bounds(with_global_expr_store_as_local(|expr_store| {
            *expr_store.get_bounds(self.get_id())
        }))
    }

    /// The absolute value. [1, F::TWO_INV[ are positives, [F::TWO_INV, -1] are negatives.
    pub fn abs(&self) -> FieldValue<F> {
        Self::new(FieldExpr::Abs(*self))
    }
    /// 0 for nonnegatives, 1 for negatives
    fn sign_bit(&self) -> FieldValue<F> {
        sign_bit(*self).into()
    }
    /// 1 if self >= 0, -1 otherwise
    pub fn sign(&self) -> FieldValue<F> {
        1 - 2 * self.sign_bit()
    }
    /// Whether the value is known to be plaintext.
    pub fn is_plaintext(&self) -> bool {
        with_global_expr_store_as_local(|expr_store| expr_store.get_is_plaintext(self.get_id()))
    }
    /// Whether the value is compile-time known.
    pub fn is_constant(&self) -> bool {
        self.as_constant().is_some()
    }
    /// If the value is compile-time known, `Some(that_value)`. Otherwise, `None`.
    pub fn as_constant(&self) -> Option<F> {
        self.bounds().as_constant()
    }
    /// Signed less than.
    pub fn signed_lt(self, other: FieldValue<F>) -> BooleanValue {
        self.param_lt(other, true)
    }
    /// Signed less or equal.
    pub fn signed_le(self, other: FieldValue<F>) -> BooleanValue {
        self.param_le(other, true)
    }
    /// Signed greater than.
    pub fn signed_gt(self, other: FieldValue<F>) -> BooleanValue {
        self.param_gt(other, true)
    }
    /// Signed greater or equal.
    pub fn signed_ge(self, other: FieldValue<F>) -> BooleanValue {
        self.param_ge(other, true)
    }
    /// Signed less than.
    pub fn param_lt(self, other: FieldValue<F>, signed: bool) -> BooleanValue {
        !BooleanValue::from(FieldValue::new(FieldExpr::Ge(self, other, signed)))
    }
    /// Signed less or equal.
    pub fn param_le(self, other: FieldValue<F>, signed: bool) -> BooleanValue {
        !BooleanValue::from(FieldValue::new(FieldExpr::Gt(self, other, signed)))
    }
    /// Signed greater than.
    pub fn param_gt(self, other: FieldValue<F>, signed: bool) -> BooleanValue {
        BooleanValue::from(FieldValue::new(FieldExpr::Gt(self, other, signed)))
    }
    /// Signed greater or equal.
    pub fn param_ge(self, other: FieldValue<F>, signed: bool) -> BooleanValue {
        BooleanValue::from(FieldValue::new(FieldExpr::Ge(self, other, signed)))
    }

    /// Random value.
    pub fn random() -> FieldValue<F> {
        let len = with_global_expr_store_as_local(|expr_store| expr_store.len());
        let res = FieldValue::new(FieldExpr::RandomVal(RandomValId::new()));
        // We are asserting that this random value is not an already existing FieldValue.
        // This bug happened before.
        assert_eq!(len, res.get_id(), "Randomness was reused !");
        res
    }
}

macro_rules! random_unsigned {
    ($t: ident, $n: expr) => {
        impl<F: ActuallyUsedField> FieldValue<F> {
            pub fn $t() -> FieldValue<F> {
                assert!($n < F::NUM_BITS);
                FieldValue::<F>::from_le_bits(
                    (0..$n)
                        .map(|_| BooleanValue::random())
                        .collect::<Vec<BooleanValue>>(),
                    false,
                )
            }
        }
    };
}

// Important: these are low-level convenience functions that can be used e.g. to generate a nonce
// (still need to reveal though). They are not intended to be exposed on the interpreter.
random_unsigned!(random_u8, 8);
random_unsigned!(random_u16, 16);
random_unsigned!(random_u32, 32);
random_unsigned!(random_u64, 64);
random_unsigned!(random_u128, 128);

impl<F: ActuallyUsedField> WithBooleanBounds for FieldValue<F> {
    fn with_boolean_bounds(&self) -> Self {
        Self::new(FieldExpr::Bounds(*self, FieldBounds::new(F::ZERO, F::ONE)))
    }
}

impl<F: ActuallyUsedField> GetBit for FieldValue<F> {
    type Output = BooleanValue;

    fn get_bit(&self, index: usize, signed: bool) -> Self::Output {
        let bounds = self.bounds();
        let circuit_size = bounds.bin_size(signed);
        let id = self.get_id();
        if circuit_size == 0 {
            BooleanValue::from(false)
        } else if index < circuit_size {
            BooleanValue::new(with_global_expr_store_as_local(|expr_store| {
                <IRBuilder as ExprStore<F>>::push_conversion(
                    expr_store,
                    BitNumToBit(id, index, signed),
                )
            }))
        } else if signed {
            BooleanValue::new(with_global_expr_store_as_local(|expr_store| {
                <IRBuilder as ExprStore<F>>::push_conversion(
                    expr_store,
                    BitNumToBit(id, circuit_size - 1, signed),
                )
            }))
        } else {
            BooleanValue::from(false)
        }
    }
}

impl<F: ActuallyUsedField> FromLeBits<BooleanValue> for FieldValue<F> {
    fn from_le_bits(bits: Vec<BooleanValue>, signed: bool) -> Self {
        FieldValue::<F>::from_id(with_global_expr_store_as_local(|expr_store| {
            <IRBuilder as ExprStore<F>>::push_conversion(
                expr_store,
                BitToBitNum(
                    bits.into_iter()
                        .map(|bit| bit.get_id())
                        .collect::<Vec<usize>>(),
                    signed,
                ),
            )
        }))
    }
}

impl<F: ActuallyUsedField> ToLeBytes for FieldValue<F> {
    type BooleanOutput = BooleanValue;

    fn to_le_bytes(self) -> [Byte<Self::BooleanOutput>; 32] {
        (0..256)
            .map(|i| self.get_bit(i, false))
            .collect::<Vec<BooleanValue>>()
            .chunks(8)
            .map(|chunk| {
                Byte::new(
                    chunk
                        .to_vec()
                        .try_into()
                        .unwrap_or_else(|v: Vec<BooleanValue>| {
                            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<F: ActuallyUsedField> Random for FieldValue<F> {
    fn random() -> Self {
        FieldValue::<F>::random()
    }
}

impl<F: ActuallyUsedField> Reveal for FieldValue<F> {
    fn reveal(self) -> Self {
        FieldValue::new(FieldExpr::Reveal(self))
    }
}

impl MxeX25519PrivateKey for FieldValue<ScalarField> {
    fn mxe_x25519_private_key() -> Self {
        Self::from_expr(Expr::Other(OtherExpr::MxeKey(MxeInput::ScalarOnly(
            MxeScalarInput::X25519PrivateKey(),
        ))))
    }
}

impl MxeRescueKey for FieldValue<BaseField> {
    fn mxe_rescue_key(i: usize) -> Self {
        assert!(i < RESCUE_KEY_COUNT);
        Self::from_expr(Expr::Other(OtherExpr::MxeKey(MxeInput::Base(
            MxeFieldInput::RescueKey(i),
        ))))
    }
}

impl MxeRescueKey for FieldValue<ScalarField> {
    fn mxe_rescue_key(i: usize) -> Self {
        assert!(i < RESCUE_KEY_COUNT);
        Self::from_expr(Expr::Other(OtherExpr::MxeKey(MxeInput::Scalar(
            MxeFieldInput::RescueKey(i),
        ))))
    }
}

macro_rules! field_value_from {
    ($t: ty) => {
        impl<F: ActuallyUsedField> From<$t> for FieldValue<F> {
            fn from(number: $t) -> FieldValue<F> {
                FieldValue::new(FieldExpr::Val(Number::from(number).into()))
            }
        }
        impl<'a, F: ActuallyUsedField> From<&'a $t> for FieldValue<F> {
            fn from(number: &'a $t) -> FieldValue<F> {
                FieldValue::new(FieldExpr::Val(Number::from(number.clone()).into()))
            }
        }
    };
}
field_value_from!(bool);
field_value_from!(i8);
field_value_from!(i16);
field_value_from!(i32);
field_value_from!(i64);
field_value_from!(i128);
field_value_from!(u8);
field_value_from!(u16);
field_value_from!(u32);
field_value_from!(u64);
field_value_from!(u128);
field_value_from!(usize);
field_value_from!(Number);

impl<F: ActuallyUsedField> From<F> for FieldValue<F> {
    fn from(number: F) -> FieldValue<F> {
        FieldValue::new(FieldExpr::Val(number))
    }
}

impl<F: ActuallyUsedField> From<BooleanValue> for FieldValue<F> {
    fn from(b: BooleanValue) -> Self {
        if b.is_plaintext() {
            FieldValue::<F>::from_id(with_global_expr_store_as_local(|expr_store| {
                expr_store.push_conversion(ScalarFromPlaintextBit::<F, usize>(b.get_id()))
            }))
        } else {
            FieldValue::<F>::from_le_bits(vec![b], false)
        }
    }
}

impl<F: ActuallyUsedField> From<Byte<BooleanValue>> for FieldValue<F> {
    fn from(byte: Byte<BooleanValue>) -> Self {
        FieldValue::<F>::from_le_bits(byte.to_vec(), false)
    }
}

impl FromF25519<FieldValue<BaseField>> for FieldValue<BaseField> {
    fn from_F25519(value: FieldValue<BaseField>) -> Vec<FieldValue<BaseField>> {
        vec![value]
    }
}

impl FromF25519<FieldValue<BaseField>> for FieldValue<ScalarField> {
    /// Injectively map a `FieldValue<BaseField>` element into a `Vec` of `FieldValue<ScalarField>`
    /// elements.
    fn from_F25519(value: FieldValue<BaseField>) -> Vec<FieldValue<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| {
                FieldValue::<ScalarField>::from_le_bits(
                    chunk
                        .iter()
                        .flat_map(|byte| byte.to_vec())
                        .collect::<Vec<BooleanValue>>(),
                    false,
                )
            })
            .collect::<Vec<FieldValue<ScalarField>>>()
    }
}

macro_rules! impl_right_scalar_binop {
    ($t: ident, $op: ident, $e: ident, $s: ty) => {
        impl<F: ActuallyUsedField> $t<$s> for FieldValue<F> {
            type Output = FieldValue<F>;
            fn $op(self, rhs: $s) -> Self::Output {
                FieldValue::new(FieldExpr::$e(self, rhs.into()))
            }
        }
        impl<'a, F: ActuallyUsedField> $t<$s> for &'a FieldValue<F> {
            type Output = FieldValue<F>;
            fn $op(self, rhs: $s) -> Self::Output {
                FieldValue::new(FieldExpr::$e(self.clone(), rhs.into()))
            }
        }
        impl<'b, F: ActuallyUsedField> $t<&'b $s> for FieldValue<F> {
            type Output = FieldValue<F>;
            fn $op(self, rhs: &'b $s) -> Self::Output {
                FieldValue::new(FieldExpr::$e(self, rhs.clone().into()))
            }
        }
        impl<'a, 'b, F: ActuallyUsedField> $t<&'b $s> for &'a FieldValue<F> {
            type Output = FieldValue<F>;
            fn $op(self, rhs: &'b $s) -> Self::Output {
                FieldValue::new(FieldExpr::$e(self.clone(), rhs.clone().into()))
            }
        }
    };
}

macro_rules! impl_right_trait_binop {
    ($t: ident, $op: ident, $e: ident) => {
        impl<F: ActuallyUsedField, T: Into<FieldValue<F>>> $t<T> for FieldValue<F> {
            type Output = FieldValue<F>;
            fn $op(self, rhs: T) -> Self::Output {
                FieldValue::new(FieldExpr::$e(self, rhs.into()))
            }
        }
    };
}

macro_rules! impl_left_scalar_binop {
    ($t: ident, $op: ident, $e: ident, $s: ty) => {
        impl<F: ActuallyUsedField> $t<FieldValue<F>> for $s {
            type Output = FieldValue<F>;
            fn $op(self, rhs: FieldValue<F>) -> Self::Output {
                FieldValue::new(FieldExpr::$e(self.into(), rhs))
            }
        }
        impl<'a, F: ActuallyUsedField> $t<FieldValue<F>> for &'a $s {
            type Output = FieldValue<F>;
            fn $op(self, rhs: FieldValue<F>) -> Self::Output {
                FieldValue::new(FieldExpr::$e(self.clone().into(), rhs))
            }
        }
        impl<'b, F: ActuallyUsedField> $t<&'b FieldValue<F>> for $s {
            type Output = FieldValue<F>;
            fn $op(self, rhs: &'b FieldValue<F>) -> Self::Output {
                FieldValue::new(FieldExpr::$e(self.into(), rhs.clone()))
            }
        }
        impl<'a, 'b, F: ActuallyUsedField> $t<&'b FieldValue<F>> for &'a $s {
            type Output = FieldValue<F>;
            fn $op(self, rhs: &'b FieldValue<F>) -> Self::Output {
                FieldValue::new(FieldExpr::$e(self.clone().into(), rhs.clone()))
            }
        }
    };
}

macro_rules! impl_sym_binop {
    ($t: ident, $op: ident, $e: ident) => {
        impl_right_trait_binop!($t, $op, $e);
        impl_left_scalar_binop!($t, $op, $e, i32);
        impl_left_scalar_binop!($t, $op, $e, Number);
        //impl_left_scalar_binop!($t, $op, $e, F);
    };
}

macro_rules! impl_bool_right_trait_binop {
    ($t: ident, $op: ident, $e: ident, $($o:tt)*) => {
        impl<F: ActuallyUsedField, T: Into<FieldValue<F>>> $t<T> for FieldValue<F> {
            type Output = BooleanValue;
            fn $op(self, rhs: T) -> Self::Output {
                BooleanValue::from(FieldValue::new(FieldExpr::$e(self, rhs.into(), $($o)*)))
            }
        }
    };
}

macro_rules! impl_bool_left_scalar_binop {
    ($t: ident, $op: ident, $e: ident, $s: ty, $($o:tt)*) => {
        impl<F: ActuallyUsedField> $t<FieldValue<F>> for $s {
            type Output = BooleanValue;
            fn $op(self, rhs: FieldValue<F>) -> Self::Output {
                BooleanValue::from(FieldValue::new(FieldExpr::$e(self.into(), rhs, $($o)*)))
            }
        }
        impl<'a, F: ActuallyUsedField> $t<FieldValue<F>> for &'a $s {
            type Output = BooleanValue;
            fn $op(self, rhs: FieldValue<F>) -> Self::Output {
                BooleanValue::from(FieldValue::new(FieldExpr::$e(self.clone().into(), rhs, $($o)*)))
            }
        }
        impl<'b, F: ActuallyUsedField> $t<&'b FieldValue<F>> for $s {
            type Output = BooleanValue;
            fn $op(self, rhs: &'b FieldValue<F>) -> Self::Output {
                BooleanValue::from(FieldValue::new(FieldExpr::$e(self.into(), rhs.clone(), $($o)*)))
            }
        }
        impl<'a, 'b, F: ActuallyUsedField> $t<&'b FieldValue<F>> for &'a $s {
            type Output = BooleanValue;
            fn $op(self, rhs: &'b FieldValue<F>) -> Self::Output {
                BooleanValue::from(FieldValue::new(FieldExpr::$e(
                    self.clone().into(),
                    rhs.clone(),
                    $($o)*
                )))
            }
        }
    };
}

macro_rules! impl_bool_sym_binop {
    ($t: ident, $op: ident, $e: ident, $($o:tt)*) => {
        impl_bool_right_trait_binop!($t, $op, $e, $($o)*);
        impl_bool_left_scalar_binop!($t, $op, $e, i32, $($o)*);
        impl_bool_left_scalar_binop!($t, $op, $e, Number, $($o)*);
        //impl_bool_left_scalar_binop!($t, $op, $e, F);
    };
}

macro_rules! impl_assign_binop {
    ($t: ident, $assign_op: ident, $op: tt) => {
        impl<F: ActuallyUsedField, T: Into<FieldValue<F>>> $t<T> for FieldValue<F> {
            fn $assign_op(&mut self, rhs: T) {
                *self = *self $op rhs;
            }
        }
    };
}

impl_sym_binop!(Add, add, Add);
impl_sym_binop!(Sub, sub, Sub);
impl_sym_binop!(Mul, mul, Mul);
impl_sym_binop!(Div, div, Div);
impl_sym_binop!(Rem, rem, Rem);
impl_bool_sym_binop!(Equal, eq, Equal,);
impl_bool_sym_binop!(GreaterEqual, ge, Ge, false);
impl_bool_sym_binop!(GreaterThan, gt, Gt, false);
impl_right_scalar_binop!(Shr, shr, LogicalRightShift, usize);
impl_assign_binop!(AddAssign, add_assign, +);
impl_assign_binop!(SubAssign, sub_assign, -);
impl_assign_binop!(MulAssign, mul_assign, *);
impl_assign_binop!(DivAssign, div_assign, /);
impl_assign_binop!(RemAssign, rem_assign, %);

impl<F: ActuallyUsedField> Neg for FieldValue<F> {
    type Output = FieldValue<F>;
    fn neg(self) -> Self::Output {
        FieldValue::new(FieldExpr::Neg(self))
    }
}
impl<F: ActuallyUsedField> Neg for &FieldValue<F> {
    type Output = FieldValue<F>;
    fn neg(self) -> Self::Output {
        FieldValue::new(FieldExpr::Neg(*self))
    }
}
impl<F: ActuallyUsedField> Not for FieldValue<F> {
    type Output = FieldValue<F>;
    fn not(self) -> Self::Output {
        1 - self
    }
}
impl<F: ActuallyUsedField> Not for &FieldValue<F> {
    type Output = FieldValue<F>;
    fn not(self) -> Self::Output {
        1 - self
    }
}
impl<F: ActuallyUsedField> Selectable<FieldValue<F>> for FieldValue<F> {
    type Conditional = FieldValue<F>;
    type Output = FieldValue<F>;

    fn construct_selection(condition: Self::Conditional, a: Self, b: Self) -> Self::Output {
        FieldValue::new(FieldExpr::Where(condition, a, b))
    }
}

impl<F: ActuallyUsedField> Selectable<FieldValue<F>> for &FieldValue<F> {
    type Conditional = FieldValue<F>;
    type Output = FieldValue<F>;

    fn construct_selection(
        condition: Self::Conditional,
        a: Self,
        b: FieldValue<F>,
    ) -> Self::Output {
        FieldValue::new(FieldExpr::Where(condition, *a, b))
    }
}
impl<F: ActuallyUsedField> Selectable<&FieldValue<F>> for FieldValue<F> {
    type Conditional = FieldValue<F>;
    type Output = FieldValue<F>;

    fn construct_selection(
        condition: Self::Conditional,
        a: Self,
        b: &FieldValue<F>,
    ) -> Self::Output {
        FieldValue::new(FieldExpr::Where(condition, a, *b))
    }
}

impl<'a, 'b, F: ActuallyUsedField> Selectable<&'b FieldValue<F>> for &'a FieldValue<F> {
    type Conditional = FieldValue<F>;
    type Output = FieldValue<F>;

    fn construct_selection(
        condition: Self::Conditional,
        a: &'a FieldValue<F>,
        b: &'b FieldValue<F>,
    ) -> Self::Output {
        FieldValue::new(FieldExpr::Where(condition, *a, *b))
    }
}

impl<F: ActuallyUsedField, T, U, V> Select<T, U, V> for FieldValue<F>
where
    T: Selectable<V, Conditional = FieldValue<F>, Output = U>,
{
    fn select(self, t: T, u: V) -> U {
        T::construct_selection(self, t, u)
    }
}
impl<F: ActuallyUsedField, T, U, V> Select<T, U, V> for &FieldValue<F>
where
    T: Selectable<V, Conditional = FieldValue<F>, Output = U>,
{
    fn select(self, t: T, u: V) -> U {
        T::construct_selection(*self, t, u)
    }
}

impl<F: ActuallyUsedField> Zero for FieldValue<F> {
    fn zero() -> Self {
        0.into()
    }

    fn is_zero(&self) -> bool {
        F::expr_to_field_expr(self.expr()) == Some(FieldExpr::Val(F::ZERO))
    }
}

impl<F: ActuallyUsedField> Sum for FieldValue<F> {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(0.into(), |acc, x| acc + x)
    }
}

impl<F: ActuallyUsedField> Invert for FieldValue<F> {
    fn invert(self, is_expected_non_zero: bool) -> Self {
        let is_zero_mask = if !is_expected_non_zero {
            FieldValue::<F>::from(self.eq(FieldValue::<F>::from(0)))
        } else {
            FieldValue::<F>::from(0)
        };
        Self::new(FieldExpr::FieldInverse(self + is_zero_mask)) - is_zero_mask
    }
}

impl<F: ActuallyUsedField> Pow for FieldValue<F> {
    fn pow(self, e: &Number, is_expected_non_zero: bool) -> Self {
        let e = e % (F::modulus() - 1);
        if are_unbuilt_circuits_allowed() {
            FieldValue::new(FieldExpr::Pow(self, e, is_expected_non_zero))
        } else {
            let pow_circuit = PowCircuit::new(e, is_expected_non_zero);
            pow_circuit.run(vec![self])[0]
        }
    }
}

// TODO: uncomment when fixed, Marius
// #[cfg(test)]
// mod tests {
//     use crate::ArcisField;
//     use super::*;
//     #[test]
//     fn random_bit() {
//         <FieldValue<ArcisField> as RandomBit>::random_bit();
//     }
// }