arcis 0.6.0-alpha

use crate::{ArcisCiphertext, ArcisType, Mxe, Shared};
use arcis_compiler::{
    traits::ToMontgomery,
    utils::{
        crypto::{key::X25519PublicKey, rescue_cipher::RescueCipher},
        curve_point::CurvePoint,
        field::{BaseField, ScalarField},
        number::Number,
        packing::{DataSize, PackLocation},
    },
    ArcisField,
    EvalValue,
};
use std::marker::PhantomData;

/// Encoded data that allows to store data on chain which can only be decoded by its owner.
/// ```
/// use arcis::*;
///
/// #[encrypted]
/// mod circuits {
///     use arcis::*;
///
///     #[instruction]
///     fn doc_enc(
///         player0_enc_u8: Enc<Shared, u8>,
///         player1_enc_u8: Enc<Shared, u8>,
///         stored_enc_u8: Enc<Mxe, u8>,
///         referee: ArcisX25519Pubkey,
///     ) -> (
///         Enc<Shared, u8>,
///         Enc<Shared, u8>,
///         Enc<Mxe, u8>,
///         Enc<Shared, Pack<[u8; 3]>>,
///     ) {
///         // We decode every `Enc` into secret-shared `u8`.
///         let player0_u8 = player0_enc_u8.to_arcis();
///         let player1_u8 = player1_enc_u8.to_arcis();
///         let stored_u8 = stored_enc_u8.to_arcis();
///         // We send player 0 what player 1 sent.
///         let player0_output = player0_enc_u8.owner.from_arcis(player1_u8);
///         // We send player 1 what was stored.
///         let player1_output = player1_enc_u8.owner.from_arcis(stored_u8);
///         // We store what player 0 sent.
///         let enc_mxe_output = stored_enc_u8.owner.from_arcis(player0_u8); // Mxe::get().from_arcis(player0_u8) works too.
///
///         // We send everything to the referee.
///         // First we pack things so that they take less place on chain.
///         let pack = Pack::new([player0_u8, player1_u8, stored_u8]);
///         let referee_output = Shared::new(referee).from_arcis(pack);
///         // Then we output every output.
///         (player0_output, player1_output, enc_mxe_output, referee_output)
///     }
/// }
/// ```
#[derive(Debug, PartialEq)]
pub struct Enc<C: Cipher, T: ArcisType> {
    #[allow(unused)]
    pub owner: C,
    pub data: EncData<T>,
}

impl<C: Cipher, T: ArcisType> Enc<C, T> {
    fn new(target: C, data: Vec<ArcisCiphertext>) -> Enc<C, T> {
        Enc {
            owner: target,
            data: EncData::new(data),
        }
    }

    /// Decode its content into secret-shares.
    pub fn to_arcis(&self) -> T {
        self.owner.to_arcis(self.data.data.as_slice())
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct EncData<T: ArcisType> {
    data: Vec<ArcisCiphertext>,
    phantom: PhantomData<T>,
}

impl<T: ArcisType> EncData<T> {
    fn new(data: Vec<ArcisCiphertext>) -> EncData<T> {
        Self {
            data,
            phantom: PhantomData,
        }
    }
}

impl<T: ArcisType> ArcisType for EncData<T> {
    fn n_values() -> usize {
        T::n_values()
    }

    fn gen_input(values: &mut Vec<EvalValue>) -> Self {
        EncData::new(
            (0..Self::n_values())
                .map(|_| BaseField::gen_input(values))
                .collect(),
        )
    }

    fn from_values(values: &[EvalValue]) -> Self {
        let data = values
            .iter()
            .map(|x| ArcisField::from(x.to_signed_number()))
            .collect();
        let phantom = PhantomData;
        Self { data, phantom }
    }

    fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
        outputs.extend(self.data.iter().cloned().map(EvalValue::Base));
    }

    fn is_similar(&self, other: &Self) -> bool {
        // could maybe be changed to `true`
        self.data == other.data
    }

    fn n_bools() -> usize {
        panic!("Cannot generate EncData.")
    }

    fn from_bools(_bools: &[bool]) -> Self {
        panic!("Cannot generate EncData.")
    }

    fn data_size(acc: &mut Vec<DataSize>) {
        acc.extend(std::iter::repeat_n(DataSize::Full, Self::n_values()));
    }

    fn pack(&self, locations: &mut &[PackLocation], containers: &mut [BaseField]) {
        for item in self.data.iter() {
            let location = locations[0];
            *locations = &locations[1..];
            assert_eq!(location.bit_offset, 0, "BaseField item has a bit offset.");
            containers[location.index] = *item;
        }
    }

    fn unpack(locations: &mut &[PackLocation], containers: &[BaseField]) -> Self {
        let v = (0..Self::n_values())
            .map(|_| {
                let location = locations[0];
                *locations = &locations[1..];
                assert_eq!(location.bit_offset, 0, "BaseField item has a bit offset.");
                containers[location.index]
            })
            .collect::<Vec<_>>();
        Self::new(v)
    }
}

/// Trait of `Shared` and `Mxe`. Gives access to `.from_arcis(data)`.
pub trait Cipher: ArcisType {
    #[doc(hidden)]
    fn encrypt_vec(&self, data: &mut Vec<ArcisField>);
    #[doc(hidden)]
    fn decrypt_vec(&self, data: &mut Vec<ArcisField>);
    #[doc(hidden)]
    fn next_cipher(self) -> Self;

    /// Encodes `data` into an `Enc<Self, T>`.
    #[allow(clippy::wrong_self_convention)]
    fn from_arcis<T: ArcisType>(self, data: T) -> Enc<Self, T>
    where
        Self: Sized,
    {
        let cipher = self.next_cipher();
        let mut data = data.to_values();
        cipher.encrypt_vec(&mut data);
        Enc::new(cipher, data)
    }

    #[doc(hidden)]
    fn to_arcis<T: ArcisType>(&self, data: &[ArcisCiphertext]) -> T {
        let mut data = data.to_vec();
        self.decrypt_vec(&mut data);
        T::from_values(&data.into_iter().map(EvalValue::Base).collect::<Vec<_>>())
    }
}

impl Cipher for Shared {
    fn encrypt_vec(&self, data: &mut Vec<ArcisField>) {
        let cipher = RescueCipher::new_with_client::<BaseField, ScalarField, CurvePoint>(
            // on plaintext points we simply set is_expected_non_identity = true
            X25519PublicKey::new(self.public_key.point, true),
        );
        *data = cipher.encrypt(data.clone(), ArcisField::from(Number::from(self.nonce)));
    }

    fn decrypt_vec(&self, data: &mut Vec<ArcisField>) {
        let cipher = RescueCipher::new_with_client::<BaseField, ScalarField, CurvePoint>(
            // on plaintext points we simply set is_expected_non_identity = true
            X25519PublicKey::new(self.public_key.point, true),
        );
        *data = cipher.decrypt(data.clone(), ArcisField::from(Number::from(self.nonce)));
    }

    fn next_cipher(self) -> Self {
        Shared {
            public_key: self.public_key,
            nonce: self.nonce + 1,
        }
    }
}

impl Cipher for Mxe {
    fn encrypt_vec(&self, data: &mut Vec<ArcisField>) {
        let cipher = RescueCipher::new_for_mxe();
        *data = cipher.encrypt(data.clone(), ArcisField::from(Number::from(self.nonce)));
    }

    fn decrypt_vec(&self, data: &mut Vec<ArcisField>) {
        let cipher = RescueCipher::new_for_mxe();
        *data = cipher.decrypt(data.clone(), ArcisField::from(Number::from(self.nonce)));
    }

    fn next_cipher(self) -> Self {
        Mxe {
            nonce: self.nonce + 1,
        }
    }
}

impl<C: Cipher, T: ArcisType> ArcisType for Enc<C, T> {
    fn n_values() -> usize {
        C::n_values() + EncData::<T>::n_values()
    }

    fn gen_input(values: &mut Vec<EvalValue>) -> Self {
        let cipher = C::gen_input(values);
        let mut data_values = Vec::new();
        let data = T::gen_input(&mut data_values);
        let mut data = data.to_values();
        cipher.encrypt_vec(&mut data);
        values.extend(data.iter().copied().map(EvalValue::Base));
        Enc::new(cipher, data)
    }

    fn from_values(mut values: &[EvalValue]) -> Self {
        let cipher = C::from_mut_values(&mut values);
        let len = T::n_values();
        let data = values[..len]
            .iter()
            .map(|x| ArcisField::from(x.to_signed_number()))
            .collect();
        Enc::new(cipher, data)
    }

    fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
        self.owner.handle_outputs(outputs);
        outputs.extend(self.data.data.iter().copied().map(EvalValue::Base));
    }

    fn is_similar(&self, other: &Self) -> bool {
        let s = self.to_arcis();
        let o = other.to_arcis();
        s.is_similar(&o)
    }

    fn n_bools() -> usize {
        panic!("Cannot generate Enc.")
    }

    fn from_bools(_bools: &[bool]) -> Self {
        panic!("Cannot generate Enc.")
    }

    fn data_size(acc: &mut Vec<DataSize>) {
        C::data_size(acc);
        EncData::<T>::data_size(acc);
    }

    fn pack(&self, locations: &mut &[PackLocation], containers: &mut [BaseField]) {
        self.owner.pack(locations, containers);
        self.data.pack(locations, containers);
    }

    fn unpack(locations: &mut &[PackLocation], containers: &[BaseField]) -> Self {
        let owner = C::unpack(locations, containers);
        let data = EncData::<T>::unpack(locations, containers);
        Self { owner, data }
    }
}

/// The Arcis x25519 pubkey.
///
/// Use `ArcisX25519Pubkey::new_from_x` to construct a new pubkey from the Montgomery x-coordinate,
/// and `pubkey.to_x()` to return the Montgomery x-coordinate.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ArcisX25519Pubkey {
    /// Internally, the key is stored as a CurvePoint, which is an element of the Ristretto group
    /// over the twisted Edwards model of the curve. For more details, see [here](https://ristretto.group/).
    point: CurvePoint,
}

impl ArcisX25519Pubkey {
    pub(crate) fn new(point: CurvePoint) -> Self {
        Self { point }
    }

    /// Construct a new ArcisX25519Pubkey from the x-coordinate of the Montgomery point.
    pub fn new_from_x(x: BaseField) -> Self {
        match X25519PublicKey::<CurvePoint>::from_le_bytes(x.to_le_bytes()) {
            Some(pubkey) => Self {
                point: pubkey.inner(),
            },
            None => Self {
                point: X25519PublicKey::<CurvePoint>::default().inner(),
            },
        }
    }

    /// Return the Montgomery x-coordinate of the ArcisX25519Pubkey.
    pub fn to_x(self) -> BaseField {
        // on plaintext points we simply set is_expected_non_identity = true
        self.point.to_montgomery(true).0
    }
}

impl ArcisType for ArcisX25519Pubkey {
    fn n_values() -> usize {
        1
    }

    fn gen_input(values: &mut Vec<EvalValue>) -> Self {
        // These values are only used for tests.
        let possible_values = [16u64, 34, 45, 49];
        use rand::Rng;

        let mut rng = rand::thread_rng();
        let choice = rng.r#gen::<usize>() % possible_values.len();
        let pick = BaseField::from(possible_values[choice]);

        let res = ArcisX25519Pubkey::new_from_x(pick);
        values.push(EvalValue::Curve(res.point));
        res
    }

    fn from_values(values: &[EvalValue]) -> Self {
        ArcisX25519Pubkey::new(values[0].to_curve())
    }

    fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
        outputs.push(EvalValue::Curve(self.point));
    }

    fn is_similar(&self, other: &Self) -> bool {
        self.point == other.point
    }

    fn n_bools() -> usize {
        panic!("Cannot generate PublicKey.")
    }

    fn from_bools(_bools: &[bool]) -> Self {
        panic!("Cannot generate PublicKey.")
    }

    fn data_size(acc: &mut Vec<DataSize>) {
        acc.push(DataSize::Full)
    }

    fn pack(&self, locations: &mut &[PackLocation], containers: &mut [BaseField]) {
        let location = locations[0];
        *locations = &locations[1..];
        assert_eq!(location.bit_offset, 0, "BaseField item has a bit offset.");
        containers[location.index] = self.to_x();
    }

    fn unpack(locations: &mut &[PackLocation], containers: &[BaseField]) -> Self {
        let location = locations[0];
        *locations = &locations[1..];
        assert_eq!(location.bit_offset, 0, "BaseField item has a bit offset.");
        Self::new_from_x(containers[location.index])
    }
}