arcis-compiler 0.9.1

use crate::{
    core::{
        actually_used_field::ActuallyUsedField,
        bounds::FieldBounds,
        circuits::traits::arithmetic_circuit::ArithmeticCircuit,
        expressions::expr::EvalFailure,
        global_value::value::FieldValue,
    },
    traits::Invert,
    utils::used_field::UsedField,
};
use std::ops::{Add, Mul, Sub};

/// Adding two affine points on the same Edwards curve
#[derive(Clone, Debug)]
pub struct AffineEdwardsPointAdditionCircuit<F: UsedField> {
    pub d: F,
}

impl<F: UsedField> AffineEdwardsPointAdditionCircuit<F> {
    #[allow(unused)]
    pub fn new(d: F) -> Self {
        Self { d }
    }
}

/// Affine point addition on the twisted Edwards curve -x^2 + y^2 = 1 + d*x^2*y^2.
#[allow(non_snake_case)]
impl<F: UsedField> AffineEdwardsPointAdditionCircuit<F> {
    pub fn add<
        T: Add<T, Output = T>
            + Sub<T, Output = T>
            + Mul<T, Output = T>
            + Invert
            + Copy
            + From<i32>
            + From<F>,
    >(
        &self,
        P1: ((T, T), bool),
        P2: ((T, T), bool),
    ) -> (T, T) {
        let ((x1, y1), is_on_curve1) = P1;
        let ((x2, y2), is_on_curve2) = P2;

        // ### round 1 ###
        // these multiplications are local if one of the two points is plaintext
        let a1 = x1 * y2;
        let a2 = y1 * x2;
        let a3 = y1 * y2;
        let a4 = x1 * x2;

        // ### round 2 ###
        let b1 = a1 * a2;

        // ### round 3 and 4 ###
        let c1 = (T::from(1) + T::from(self.d) * b1).invert(is_on_curve1 && is_on_curve2);
        let c2 = (T::from(1) - T::from(self.d) * b1).invert(is_on_curve1 && is_on_curve2);

        // ### round 5 ###
        let x_sum = (a1 + a2) * c1;
        let y_sum = (a3 + a4) * c2;

        (x_sum, y_sum)
    }
}

impl<F: UsedField> ArithmeticCircuit<F> for AffineEdwardsPointAdditionCircuit<F> {
    fn eval(&self, x: Vec<F>) -> Result<Vec<F>, EvalFailure> {
        if x.len() != 4 {
            panic!("Affine Edwards Point Addition requires input of length 4");
        }
        let (x_sum, y_sum) = Self::add(self, ((x[0], x[1]), false), ((x[2], x[3]), false));
        Ok(vec![x_sum, y_sum])
    }

    fn bounds(&self, _bounds: Vec<FieldBounds<F>>) -> Vec<FieldBounds<F>> {
        vec![FieldBounds::All, FieldBounds::All]
    }

    fn run(&self, vals: Vec<FieldValue<F>>) -> Vec<FieldValue<F>>
    where
        F: ActuallyUsedField,
    {
        if vals.len() != 4 {
            panic!("Affine Edwards Point Addition requires input of length 4");
        }
        let (x_sum, y_sum) = Self::add(
            self,
            ((vals[0], vals[1]), false),
            ((vals[2], vals[3]), false),
        );
        vec![x_sum, y_sum]
    }
}

/// Adding two projective points on the same Edwards curve
#[derive(Clone, Debug)]
pub struct ProjectiveEdwardsPointAdditionCircuit<F: UsedField> {
    pub d: F,
}

impl<F: UsedField> ProjectiveEdwardsPointAdditionCircuit<F> {
    #[allow(unused)]
    pub fn new(d: F) -> Self {
        Self { d }
    }
}

/// Projective point addition on the twisted Edwards curve -X^2*Z^2 + Y^2*Z^2 = Z^4 + d*X^2*Y^2.
#[allow(non_snake_case)]
impl<F: UsedField> ProjectiveEdwardsPointAdditionCircuit<F> {
    pub fn add<
        T: Add<T, Output = T> + Sub<T, Output = T> + Mul<T, Output = T> + Copy + From<i32> + From<F>,
    >(
        &self,
        P1: (T, T, T),
        P2: (T, T, T),
    ) -> (T, T, T) {
        let (X1, Y1, Z1) = P1;
        let (X2, Y2, Z2) = P2;

        // ### round 1 ###
        // these multiplications are local if one of the two points is plaintext
        let a1 = X1 * Y2;
        let a2 = Y1 * X2;
        let a3 = Y1 * Y2;
        let a4 = X1 * X2;
        let a5 = Z1 * Z2;

        // ### round 2 ###
        let b1 = (a1 + a2) * a5;
        let b2 = (a3 + a4) * a5;
        let b3 = a5 * a5;
        let b4 = a1 * a2;

        // ### round 3 ###
        let X_sum = b1 * (b3 - T::from(self.d) * b4);
        let Y_sum = b2 * (b3 + T::from(self.d) * b4);
        let Z_sum = (b3 + T::from(self.d) * b4) * (b3 - T::from(self.d) * b4);

        (X_sum, Y_sum, Z_sum)
    }
}

impl<F: UsedField> ArithmeticCircuit<F> for ProjectiveEdwardsPointAdditionCircuit<F> {
    #[allow(non_snake_case)]
    fn eval(&self, x: Vec<F>) -> Result<Vec<F>, EvalFailure> {
        if x.len() != 6 {
            panic!("Projective Edwards Point Addition requires input of length 6");
        }
        let (X_sum, Y_sum, Z_sum) = Self::add(self, (x[0], x[1], x[2]), (x[3], x[4], x[5]));
        Ok(vec![X_sum, Y_sum, Z_sum])
    }

    fn bounds(&self, _bounds: Vec<FieldBounds<F>>) -> Vec<FieldBounds<F>> {
        vec![FieldBounds::All; 3]
    }

    #[allow(non_snake_case)]
    fn run(&self, vals: Vec<FieldValue<F>>) -> Vec<FieldValue<F>>
    where
        F: ActuallyUsedField,
    {
        if vals.len() != 6 {
            panic!("Projective Edwards Point Addition requires input of length 6");
        }
        let (X_sum, Y_sum, Z_sum) = Self::add(
            self,
            (vals[0], vals[1], vals[2]),
            (vals[3], vals[4], vals[5]),
        );
        vec![X_sum, Y_sum, Z_sum]
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        core::circuits::traits::arithmetic_circuit::tests::TestedArithmeticCircuit,
        traits::FromLeBytes,
        utils::{
            elliptic_curve::{ProjectiveEdwardsPoint, EDWARDS25519_D},
            field::BaseField,
        },
    };
    use ff::Field;
    use rand::Rng;

    /// Generate a random affine point on the twisted Edwards curve -x^2 + y^2 = 1 + d*x^2*y^2.
    fn gen_random_affine_point<R: Rng + ?Sized>(
        rng: &mut R,
        d: BaseField,
    ) -> (BaseField, BaseField) {
        fn gen_x_y<R: Rng + ?Sized>(rng: &mut R, d: BaseField) -> Option<(BaseField, BaseField)> {
            let x = BaseField::random(&mut *rng);
            let x2 = x * x;
            let a = BaseField::ONE + x2;
            let b = BaseField::ONE - d * x2;
            let y2 = a * b.invert(true);
            let sqrt = y2.sqrt();
            if sqrt.is_some().into() {
                Some((x, sqrt.unwrap()))
            } else {
                None
            }
        }
        let mut xy = gen_x_y(rng, d);
        while xy.is_none() {
            xy = gen_x_y(rng, d)
        }
        let (x, y) = xy.unwrap();
        let x2 = x * x;
        let y2 = y * y;
        assert_eq!(-x2 + y2, BaseField::ONE + d * x2 * y2);
        (x, y)
    }

    /// Generate a random projective point on the twisted Edwards curve -X^2*Z^2 + Y^2*Z^2 = Z^4 +
    /// d*X^2*Y^2.
    fn gen_random_projective_point<R: Rng + ?Sized>(
        rng: &mut R,
        d: BaseField,
    ) -> (BaseField, BaseField, BaseField) {
        let (x, y) = gen_random_affine_point(rng, d);
        (x, y, BaseField::ONE)
    }

    #[test]
    #[allow(non_snake_case)]
    fn test_affine_add() {
        let addition_circuit = AffineEdwardsPointAdditionCircuit::<BaseField>::new(
            BaseField::from_le_bytes(EDWARDS25519_D),
        );
        let O = (BaseField::ZERO, BaseField::ONE);
        // the unique Fp-rational 2-torsion point
        let Q2 = (BaseField::ZERO, -BaseField::ONE);
        let two_Q2 = addition_circuit.add((Q2, true), (Q2, true));
        assert_eq!(two_Q2, O);
        let rng = &mut crate::utils::test_rng::get();
        for _ in 0..64 {
            let P = gen_random_affine_point(rng, BaseField::from_le_bytes(EDWARDS25519_D));
            let P_plus_O = addition_circuit.add((P, true), (O, true));
            let O_plus_P = addition_circuit.add((O, true), (P, true));
            assert_eq!(P_plus_O, P);
            assert_eq!(O_plus_P, P);
            let neg_P = (-P.0, P.1);
            let P_minus_P = addition_circuit.add((P, true), (neg_P, true));
            let neg_P_plus_P = addition_circuit.add((neg_P, true), (P, true));
            assert_eq!(P_minus_P, O);
            assert_eq!(neg_P_plus_P, O);
        }
    }

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

        fn gen_n_inputs<R: Rng + ?Sized>(&self, _rng: &mut R) -> usize {
            4
        }
    }
    #[test]
    fn tested_affine() {
        AffineEdwardsPointAdditionCircuit::<BaseField>::test(1, 16)
    }

    #[test]
    #[allow(non_snake_case)]
    fn test_projective_add() {
        let addition_circuit = ProjectiveEdwardsPointAdditionCircuit::<BaseField>::new(
            BaseField::from_le_bytes(EDWARDS25519_D),
        );
        let O = (BaseField::ZERO, BaseField::ONE, BaseField::ONE);
        // the unique Fp-rational 2-torsion point
        let Q2 = (BaseField::ZERO, -BaseField::ONE, BaseField::ONE);
        let two_Q2 = addition_circuit.add(Q2, Q2);
        assert_eq!(two_Q2, O);
        let rng = &mut crate::utils::test_rng::get();
        for _ in 0..64 {
            let P = gen_random_projective_point(rng, BaseField::from_le_bytes(EDWARDS25519_D));
            let P_plus_O = addition_circuit.add(P, O);
            let O_plus_P = addition_circuit.add(O, P);
            assert_eq!(P_plus_O, P);
            assert_eq!(O_plus_P, P);
            let neg_P = (-P.0, P.1, P.2);
            let P_minus_P = addition_circuit.add(P, neg_P);
            let neg_P_plus_P = addition_circuit.add(neg_P, P);
            assert_eq!(
                ProjectiveEdwardsPoint::new(P_minus_P, true, true),
                ProjectiveEdwardsPoint::new(O, true, true)
            );
            assert_eq!(
                ProjectiveEdwardsPoint::new(neg_P_plus_P, true, true),
                ProjectiveEdwardsPoint::new(O, true, true)
            );
        }
    }

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

        fn gen_n_inputs<R: Rng + ?Sized>(&self, _rng: &mut R) -> usize {
            6
        }
    }
    #[test]
    fn tested_pojective() {
        ProjectiveEdwardsPointAdditionCircuit::<BaseField>::test(1, 64)
    }
}