arcis-compiler 0.9.7

use crate::{
    core::{
        actually_used_field::ActuallyUsedField,
        bounds::{BoolBounds, Bounds, CurveBounds, FieldBounds},
        circuits::boolean::byte::Byte,
        expressions::expr::EvalValue,
    },
    traits::FromLeBytes,
    utils::{
        crypto::key::{
            AES_128_KEY_COUNT,
            AES_192_KEY_COUNT,
            AES_256_KEY_COUNT,
            ED25519_SECRET_KEY_COUNT,
            ED25519_SIGNING_KEY_HASH_PREFIX_COUNT,
            ED25519_VERIFYING_KEY_COUNT,
            MXE_AES128_KEY,
            MXE_AES192_KEY,
            MXE_AES256_KEY,
            MXE_ED25519_SECRET_KEY,
            MXE_ED25519_SIGNING_KEY_HASH_PREFIX,
            MXE_ED25519_SIGNING_KEY_S,
            MXE_ED25519_VERIFYING_KEY,
            MXE_ELGAMAL_PUBKEY,
            MXE_ELGAMAL_SECRET_KEY,
            MXE_RESCUE_BASE_FIELD_KEY,
            MXE_X25519_PRIVATE_KEY,
            RESCUE_KEY_COUNT,
        },
        curve_point::CurvePoint,
        field::{BaseField, ScalarField},
    },
};
use curve25519_dalek_arcium_fork::ristretto::CompressedRistretto;
use serde::{Deserialize, Serialize};
use std::marker::PhantomData;

/// `pub trait` but meant to be closed.
pub trait ArxInput {
    type Output;
    /// Whether this `MxeInput` is plaintext.
    fn is_plaintext(&self) -> bool;
    /// The mock value of this `MxeInput`.
    fn mock_eval(&self) -> Self::Output;
}

/// An input given by the MXE, not by the user.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum MxeInput {
    Bit(MxeBitInput),
    ScalarOnly(MxeScalarInput),
    Scalar(MxeFieldInput<ScalarField>),
    Base(MxeFieldInput<BaseField>),
    Curve(MxeCurveInput),
}

impl MxeInput {
    /// The bounds of this `MxeInput`.
    pub fn bounds(&self) -> Bounds {
        match self {
            MxeInput::Bit(_) => Bounds::Bit(BoolBounds::new(true, true)),
            MxeInput::ScalarOnly(_) => Bounds::Scalar(FieldBounds::All),
            MxeInput::Scalar(i) => Bounds::Scalar(i.bounds()),
            MxeInput::Base(i) => Bounds::Base(i.bounds()),
            MxeInput::Curve(_) => Bounds::Curve(CurveBounds::All),
        }
    }
}

impl ArxInput for MxeInput {
    type Output = EvalValue;
    fn is_plaintext(&self) -> bool {
        match self {
            MxeInput::Bit(i) => i.is_plaintext(),
            MxeInput::ScalarOnly(i) => i.is_plaintext(),
            MxeInput::Scalar(i) => i.is_plaintext(),
            MxeInput::Base(i) => i.is_plaintext(),
            MxeInput::Curve(i) => i.is_plaintext(),
        }
    }
    fn mock_eval(&self) -> Self::Output {
        match self {
            MxeInput::Bit(i) => EvalValue::Bit(i.mock_eval()),
            MxeInput::ScalarOnly(i) => EvalValue::Scalar(i.mock_eval()),
            MxeInput::Scalar(i) => EvalValue::Scalar(i.mock_eval()),
            MxeInput::Base(i) => EvalValue::Base(i.mock_eval()),
            MxeInput::Curve(i) => EvalValue::Curve(i.mock_eval()),
        }
    }
}

/// A `bool` MXE input.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum MxeBitInput {
    /// The MXE aes-128 key.
    AES128Key(usize),
    /// The MXE aes-192 key.
    AES192Key(usize),
    /// The MXE aes-256 key.
    AES256Key(usize),
    /// The MXE ed25519 secret key.
    Ed25519SecretKey(usize),
    /// The MXE ed25519 signing key hash-prefix.
    Ed25519SigningKeyHashPrefix(usize),
}

impl ArxInput for MxeBitInput {
    type Output = bool;

    fn is_plaintext(&self) -> bool {
        false
    }

    fn mock_eval(&self) -> Self::Output {
        match *self {
            MxeBitInput::AES128Key(i) => {
                if i < AES_128_KEY_COUNT {
                    Byte::from(MXE_AES128_KEY[i / 8]).get_bits()[i % 8]
                } else {
                    false
                }
            }
            MxeBitInput::AES192Key(i) => {
                if i < AES_192_KEY_COUNT {
                    Byte::from(MXE_AES192_KEY[i / 8]).get_bits()[i % 8]
                } else {
                    false
                }
            }
            MxeBitInput::AES256Key(i) => {
                if i < AES_256_KEY_COUNT {
                    Byte::from(MXE_AES256_KEY[i / 8]).get_bits()[i % 8]
                } else {
                    false
                }
            }
            MxeBitInput::Ed25519SecretKey(i) => {
                if i < ED25519_SECRET_KEY_COUNT {
                    Byte::from(MXE_ED25519_SECRET_KEY[i / 8]).get_bits()[i % 8]
                } else {
                    false
                }
            }
            MxeBitInput::Ed25519SigningKeyHashPrefix(i) => {
                if i < ED25519_SIGNING_KEY_HASH_PREFIX_COUNT {
                    Byte::from(MXE_ED25519_SIGNING_KEY_HASH_PREFIX[i / 8]).get_bits()[i % 8]
                } else {
                    false
                }
            }
        }
    }
}

/// A `ScalarField` MXE input.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum MxeScalarInput {
    /// The MXE x25519 private key.
    X25519PrivateKey(),
    /// The MXE ed25519 signing key s.
    Ed25519SigningKeyS(),
    /// The ElGamal secret key, used for c-spl.
    ElGamalSecretKey(),
}

impl ArxInput for MxeScalarInput {
    type Output = ScalarField;

    fn is_plaintext(&self) -> bool {
        false
    }

    fn mock_eval(&self) -> Self::Output {
        match self {
            MxeScalarInput::X25519PrivateKey() => {
                ScalarField::from_le_bytes(MXE_X25519_PRIVATE_KEY)
            }
            MxeScalarInput::Ed25519SigningKeyS() => {
                ScalarField::from_le_bytes(MXE_ED25519_SIGNING_KEY_S)
            }
            MxeScalarInput::ElGamalSecretKey() => {
                ScalarField::from_le_bytes(MXE_ELGAMAL_SECRET_KEY)
            }
        }
    }
}

/// An MXE input that could exist on any field.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum MxeFieldInput<F: ActuallyUsedField> {
    /// The MXE ed25519 verifying key.
    Ed25519VerifyingKey(usize),
    /// Rescue Key.
    RescueKey(usize),
    /// Do not use.
    Phantom(PhantomData<F>),
}

impl<F: ActuallyUsedField> MxeFieldInput<F> {
    pub fn bounds(&self) -> FieldBounds<F> {
        match self {
            MxeFieldInput::Ed25519VerifyingKey(_) => FieldBounds::new(F::ZERO, F::from(255)),
            MxeFieldInput::RescueKey(_) => FieldBounds::All,
            MxeFieldInput::Phantom(_) => {
                panic!("Phantom type cannot be created.")
            }
        }
    }
}

impl<F: ActuallyUsedField> ArxInput for MxeFieldInput<F> {
    type Output = F;
    fn is_plaintext(&self) -> bool {
        match self {
            MxeFieldInput::Ed25519VerifyingKey(_) => true,
            MxeFieldInput::RescueKey(_) => false,
            MxeFieldInput::Phantom(_) => {
                panic!("Phantom type is not supported")
            }
        }
    }
    fn mock_eval(&self) -> F {
        match *self {
            MxeFieldInput::Ed25519VerifyingKey(i) => {
                if i < ED25519_VERIFYING_KEY_COUNT {
                    F::from(MXE_ED25519_VERIFYING_KEY[i] as u64)
                } else {
                    F::ZERO
                }
            }
            MxeFieldInput::RescueKey(i) => {
                if i < RESCUE_KEY_COUNT {
                    F::from_le_bytes(MXE_RESCUE_BASE_FIELD_KEY[i])
                } else {
                    F::ZERO
                }
            }
            MxeFieldInput::Phantom(_) => {
                panic!("Phantom type is not supported")
            }
        }
    }
}

/// An MXE input that is a curve point.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum MxeCurveInput {
    /// The ElGamal pubkey, used for c-spl.
    ElGamalPubkey(),
}

impl ArxInput for MxeCurveInput {
    type Output = CurvePoint;

    fn is_plaintext(&self) -> bool {
        true
    }
    fn mock_eval(&self) -> CurvePoint {
        match self {
            MxeCurveInput::ElGamalPubkey() => {
                let point = CompressedRistretto::from_slice(&MXE_ELGAMAL_PUBKEY)
                    .unwrap()
                    .decompress()
                    .unwrap();
                CurvePoint::new(point)
            }
        }
    }
}