arcis-compiler 0.9.4

use crate::{
    core::{
        bounds::FieldBounds,
        circuits::{
            boolean::{
                boolean_value::{Boolean, BooleanValue},
                byte::Byte,
            },
            traits::arithmetic_circuit::ArithmeticCircuit,
        },
        expressions::expr::EvalFailure,
        global_value::value::FieldValue,
    },
    traits::{Equal, FromLeBytes, GetBit, GreaterEqual, Select},
    utils::{
        elliptic_curve::{
            AffineEdwardsPoint,
            EDWARDS25519_D,
            EIGHT_INV_MOD_ELL,
            F25519,
            SQRT_NEG_ONE,
        },
        field::BaseField,
        used_field::UsedField,
    },
};
use ff::Field;
use std::ops::Shr;

/// Encodes a base field element into a point on Edwards25519 (in the large prime order
/// subgroup). To make the decoding work we require the base field element to be at most 2^128.
#[derive(Clone, Debug)]
#[allow(dead_code)]
pub struct BaseFieldToEdwards25519EncodingCircuit {
    d: BaseField,
    kappa: i32,
}

impl Default for BaseFieldToEdwards25519EncodingCircuit {
    fn default() -> Self {
        Self::new()
    }
}

impl BaseFieldToEdwards25519EncodingCircuit {
    #[allow(dead_code)]
    pub fn new() -> Self {
        Self {
            d: BaseField::from_le_bytes(EDWARDS25519_D),
            kappa: 64,
        }
    }

    #[cfg(test)]
    pub fn get_d(&self) -> BaseField {
        self.d
    }
}

impl BaseFieldToEdwards25519EncodingCircuit {
    /// Encodes the message m into the y-coordinate of a point on Edwards25519.
    fn encode<
        B: Boolean + Select<T, T, T>,
        T: F25519 + GetBit<Output = B> + Equal<T, Output = B> + From<B>,
    >(
        &self,
        m: T,
    ) -> AffineEdwardsPoint<T> {
        // TODO: can we add the requirement on the input being at most 2^128 here?

        // we need to add an offset (as y = 0 and y = 1 in the below lead to a 4-torsion point or
        // the identity respectively) TODO: maybe we can omit the offset, hence m = 0 gets encoded
        // as the identity after clearing the co-factor
        let m_offset = m + T::from(2);
        // TODO: the below should be batched
        let mut points = (0..self.kappa)
            .map(|i| {
                let y = m_offset * T::from(self.kappa) + T::from(i);
                let (is_on_curve, point) = AffineEdwardsPoint::<T>::try_from_y(y);
                (is_on_curve, point.x, point.y)
            })
            .collect::<Vec<(B, T, T)>>();

        // TODO: implement in depth log2(kappa)
        let (x, y) = points.split_off(1).into_iter().fold(
            (points[0].1, points[0].2),
            |(x_current, y_current), (is_on_curve_cand, x_cand, y_cand)| {
                let x = is_on_curve_cand.select(x_cand, x_current);
                let y = is_on_curve_cand.select(y_cand, y_current);
                (x, y)
            },
        );

        // TODO: here, we will switch from base field secret-sharing to EC point secret-sharing.
        // When generating a random curve point Lambda in the prime order subgroup
        // (both in EC point and base field secret-sharing) one must not forget to add
        // a 8-torsion part to the base field secret-shared mask.

        AffineEdwardsPoint::new((x, y), true, false)
            .to_projective()
            .mul_str("0001")
            .to_affine()
    }
}

impl ArithmeticCircuit<BaseField> for BaseFieldToEdwards25519EncodingCircuit {
    fn eval(&self, x: Vec<BaseField>) -> Result<Vec<BaseField>, EvalFailure> {
        if x.len() != 1 {
            panic!("Base field to Edwards25519 encoding requires input of length 1");
        }
        if x[0] > BaseField::power_of_two(128) {
            EvalFailure::err_ub("Input must be at most 2^128")
        } else {
            let (x, y) = Self::encode::<bool, _>(self, x[0]).inner();
            Ok(vec![x, y])
        }
    }

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

    fn run(&self, vals: Vec<FieldValue<BaseField>>) -> Vec<FieldValue<BaseField>> {
        if vals.len() != 1 {
            panic!("Base field to Edwards25519 encoding requires input of length 1");
        }
        // TODO: add some diagnostics
        // let bounds = BaseField::bounds_to_field_bounds(vals[0].bounds());
        // if bounds.has_negatives() || bounds.unsigned_max() > BaseField::power_of_two(128) {
        //     EvalFailure::ub("Input must be at most 2^128")
        // }
        let (x, y) = Self::encode::<BooleanValue, _>(self, vals[0]).inner();
        vec![x, y]
    }
}

/// Decodes a base field element which was encoded via BaseFieldToEdwards25519EncodingCircuit.
#[derive(Clone, Debug)]
#[allow(dead_code)]
pub struct Edwards25519ToBaseFieldDecodingCircuit {
    d: BaseField,
    kappa: i32,
}

impl Default for Edwards25519ToBaseFieldDecodingCircuit {
    fn default() -> Self {
        Self::new()
    }
}

impl Edwards25519ToBaseFieldDecodingCircuit {
    #[allow(dead_code)]
    pub fn new() -> Self {
        Self {
            d: BaseField::from_le_bytes(EDWARDS25519_D),
            kappa: 64,
        }
    }
}

#[allow(non_snake_case, dead_code)]
impl Edwards25519ToBaseFieldDecodingCircuit {
    /// Computes the minimum y-coordinate of the 8-torsion coset of a point P.
    /// Note that Q is any of the four rational points of order 8.
    fn recover_message<
        B: Boolean + Select<T, T, T>,
        T: F25519 + Shr<usize, Output = T> + GreaterEqual<T, Output = B>,
    >(
        &self,
        P: AffineEdwardsPoint<T>,
        P_plus_Q: AffineEdwardsPoint<T>,
    ) -> T {
        let sqrt_neg_one = T::from(BaseField::from_le_bytes(SQRT_NEG_ONE));
        // Note: the action of the rational 4-torsion on a point (x, y) yields the coset {(x, y),
        // (sqrt(-1)*y, sqrt(-1)*x), (-x, -y), (-sqrt(-1)*y, -sqrt(-1)*x)}. That is, the set of
        // y-coordinates of that coset is precisely {y, -y, sqrt(-1)*x, -sqrt(-1)*x}.
        let mut y_coordinates = vec![
            P.y,
            -P.y,
            sqrt_neg_one * P.x,
            -sqrt_neg_one * P.x,
            P_plus_Q.y,
            -P_plus_Q.y,
            sqrt_neg_one * P_plus_Q.x,
            -sqrt_neg_one * P_plus_Q.x,
        ];

        // Compute the minimum y-coordinate and subtract the offset.
        y_coordinates.split_off(1).into_iter().fold(
            y_coordinates[0] >> (self.kappa.ilog2() as usize),
            |current, cand| {
                let cand = cand >> (self.kappa.ilog2() as usize);
                (cand.lt(current)).select(cand, current)
            },
        ) - T::from(2)
    }

    /// Decodes a base field element which was encoded via BaseFieldToEdwards25519EncodingCircuit.
    fn decode<
        B: Boolean + Select<T, T, T>,
        T: F25519 + Shr<usize, Output = T> + GreaterEqual<T, Output = B>,
    >(
        &self,
        P: AffineEdwardsPoint<T>,
    ) -> T {
        let P_prime = P
            .to_projective()
            .mul_bits(
                EIGHT_INV_MOD_ELL
                    .into_iter()
                    .flat_map(|byte| Byte::<bool>::from(byte).to_vec())
                    .collect::<Vec<bool>>(),
            )
            .to_affine();
        let P_prime_plus_Q = P_prime + AffineEdwardsPoint::eight_torsion_point();
        Self::recover_message(self, P_prime, P_prime_plus_Q)
    }
}

impl ArithmeticCircuit<BaseField> for Edwards25519ToBaseFieldDecodingCircuit {
    #[allow(non_snake_case)]
    fn eval(&self, x: Vec<BaseField>) -> Result<Vec<BaseField>, EvalFailure> {
        if x.len() != 2 {
            panic!("Edwards25519 to Base field decoding requires input of length 2");
        }

        let P = AffineEdwardsPoint::new((x[0], x[1]), false, false);
        let m = Self::decode::<_, _>(self, P);
        if m > BaseField::power_of_two(128) {
            EvalFailure::err_ub("Circuit only supported for decoded messages of at most 2^128")
        } else {
            Ok(vec![m])
        }
    }

    fn bounds(&self, _bounds: Vec<FieldBounds<BaseField>>) -> Vec<FieldBounds<BaseField>> {
        vec![FieldBounds::new(
            BaseField::ZERO,
            BaseField::power_of_two(128),
        )]
    }

    #[allow(non_snake_case)]
    fn run(&self, vals: Vec<FieldValue<BaseField>>) -> Vec<FieldValue<BaseField>> {
        if vals.len() != 2 {
            panic!("Edwards25519 to Base field decoding requires input of length 2");
        }

        let P = AffineEdwardsPoint::new((vals[0], vals[1]), false, false);
        vec![Self::decode::<_, _>(self, P)]
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::circuits::traits::arithmetic_circuit::tests::TestedArithmeticCircuit;
    use rand::Rng;

    #[test]
    #[allow(non_snake_case)]
    fn test_encoding_decoding() {
        let rng = &mut crate::utils::test_rng::get();
        let encoding_circuit = BaseFieldToEdwards25519EncodingCircuit::new();
        let decoding_circuit = Edwards25519ToBaseFieldDecodingCircuit::new();
        for _ in 0..2 {
            for magnitude in [8, 128] {
                let m = BaseField::gen_inclusive_range(
                    &mut *rng,
                    BaseField::ZERO,
                    BaseField::power_of_two(magnitude),
                );
                let P = encoding_circuit.encode::<bool, _>(m);
                let (x, y) = P.inner();
                let x2 = x * x;
                let y2 = y * y;
                assert_eq!(
                    -x2 + y2,
                    BaseField::ONE + encoding_circuit.get_d() * x2 * y2
                );
                let m_dec = decoding_circuit.decode::<_, _>(P);
                assert_eq!(m_dec, m);
            }
        }
    }

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

        fn gen_n_inputs<R: Rng + ?Sized>(&self, _rng: &mut R) -> usize {
            1
        }

        fn extra_checks(&self, inputs: Vec<BaseField>, outputs: Vec<BaseField>) {
            let decoding_circuit = Edwards25519ToBaseFieldDecodingCircuit::new();
            let decoded = decoding_circuit.eval(outputs).unwrap();
            assert_eq!(inputs, decoded);
        }
    }

    #[test]
    fn tested_base_field_to_edwards25519_encoding() {
        BaseFieldToEdwards25519EncodingCircuit::test(1, 4)
    }
}