arcis-compiler 0.9.7

//! Arcis implementation of <https://github.com/solana-program/zk-elgamal-proof/blob/main/zk-sdk/src/encryption/elgamal.rs>

use crate::{
    core::{
        expressions::{expr::Expr, other_expr::OtherExpr},
        global_value::{curve_value::CurveValue, value::FieldValue},
        mxe_input::{MxeInput, MxeScalarInput},
    },
    traits::{Invert, Reveal},
    utils::{
        field::ScalarField,
        zkp::pedersen::{Pedersen, PedersenCommitment, PedersenOpening},
    },
    MxeCurveInput,
};
use std::{
    ops::{Add, Mul, Sub},
    sync::LazyLock,
};
use zk_elgamal_proof::encryption::pedersen::H;

/// Algorithm handle for the twisted ElGamal encryption scheme
pub struct ElGamal;
impl ElGamal {
    /// Generates an ElGamal keypair.
    ///
    /// This function is randomized. It internally samples a scalar element using a random singlet.
    fn keygen() -> ElGamalKeypair {
        // secret scalar should be non-zero except with negligible probability
        let s = FieldValue::<ScalarField>::random();
        Self::keygen_with_scalar(&s)
    }

    /// Generates an ElGamal keypair from a scalar input that determines the ElGamal private key.
    fn keygen_with_scalar(s: &FieldValue<ScalarField>) -> ElGamalKeypair {
        let secret = ElGamalSecretKey(*s);
        let public = ElGamalPubkey::new(&secret);

        ElGamalKeypair { public, secret }
    }

    /// On input an ElGamal public key and an amount to be encrypted, the function returns a
    /// corresponding ElGamal ciphertext.
    ///
    /// This function is randomized. It internally samples a scalar element using a random singlet.
    fn encrypt(public: &ElGamalPubkey, amount: FieldValue<ScalarField>) -> ElGamalCiphertext {
        let (commitment, opening) = Pedersen::new(amount);
        let handle = public.decrypt_handle(&opening);

        ElGamalCiphertext { commitment, handle }
    }

    /// On input a public key, amount, and Pedersen opening, the function returns the corresponding
    /// ElGamal ciphertext.
    fn encrypt_with(
        amount: FieldValue<ScalarField>,
        public: &ElGamalPubkey,
        opening: &PedersenOpening,
    ) -> ElGamalCiphertext {
        let commitment = Pedersen::with(amount, opening);
        let handle = public.decrypt_handle(opening);

        ElGamalCiphertext { commitment, handle }
    }

    /// On input an amount, the function returns a twisted ElGamal ciphertext where the associated
    /// Pedersen opening is always zero. Since the opening is zero, any twisted ElGamal ciphertext
    /// of this form is a valid ciphertext under any ElGamal public key.
    pub fn encode(amount: FieldValue<ScalarField>) -> ElGamalCiphertext {
        let commitment = Pedersen::encode(amount);
        let handle = DecryptHandle(CurveValue::identity());

        ElGamalCiphertext { commitment, handle }
    }

    /// On input a secret key and a ciphertext, the function returns the encoding of the original
    /// amount.
    fn decrypt(secret: &ElGamalSecretKey, ciphertext: &ElGamalCiphertext) -> CurveValue {
        *ciphertext.commitment.get_point() - secret.0 * ciphertext.handle.0
    }
}

/// A (twisted) ElGamal encryption keypair.
#[derive(Clone, Copy)]
pub struct ElGamalKeypair {
    /// The public half of this keypair.
    public: ElGamalPubkey,
    /// The secret half of this keypair.
    secret: ElGamalSecretKey,
}

impl ElGamalKeypair {
    /// Convert an ElGamal secret key to an ElGamal keypair.
    pub fn new(secret: ElGamalSecretKey) -> Self {
        let public = ElGamalPubkey::new(&secret);
        Self { public, secret }
    }

    /// Generates the public and secret keys for ElGamal encryption.
    ///
    /// This function is randomized. It internally samples a scalar element using a random singlet.
    pub fn new_rand() -> Self {
        ElGamal::keygen()
    }

    /// Create an ElGamal keypair from an ElGamal public key and an ElGamal secret key.
    pub fn new_from_inner(public: ElGamalPubkey, secret: ElGamalSecretKey) -> Self {
        Self { public, secret }
    }

    /// Loads the MXE ElGamal keypair.
    pub fn mxe_keypair() -> Self {
        Self {
            public: ElGamalPubkey(CurveValue::from_expr(Expr::Other(OtherExpr::MxeKey(
                MxeInput::Curve(MxeCurveInput::ElGamalPubkey()),
            )))),
            secret: ElGamalSecretKey(FieldValue::<ScalarField>::from_expr(Expr::Other(
                OtherExpr::MxeKey(MxeInput::ScalarOnly(MxeScalarInput::ElGamalSecretKey())),
            ))),
        }
    }

    pub fn pubkey(&self) -> &ElGamalPubkey {
        &self.public
    }

    pub fn secret(&self) -> &ElGamalSecretKey {
        &self.secret
    }
}

/// Public key for the ElGamal encryption scheme.
#[derive(Clone, Copy, Default)]
pub struct ElGamalPubkey(CurveValue);
impl ElGamalPubkey {
    /// Derives the `ElGamalPubkey` that uniquely corresponds to an `ElGamalSecretKey`.
    pub fn new(secret: &ElGamalSecretKey) -> Self {
        let s = &secret.0;

        ElGamalPubkey((s.invert(true) * CurveValue::from(*LazyLock::force(&H))).reveal())
    }

    /// Create an `ElGamalPubkey` from a `CurveValue`.
    pub fn new_from_inner(public: CurveValue) -> Self {
        assert!(public.is_plaintext());
        Self(public)
    }

    pub fn get_point(&self) -> &CurveValue {
        &self.0
    }

    /// Encrypts an amount under the public key.
    ///
    /// This function is randomized. It internally samples a scalar element using a random singlet.
    pub fn encrypt(&self, amount: FieldValue<ScalarField>) -> ElGamalCiphertext {
        ElGamal::encrypt(self, amount)
    }

    /// Encrypts an amount under the public key and an input Pedersen opening.
    pub fn encrypt_with(
        &self,
        amount: FieldValue<ScalarField>,
        opening: &PedersenOpening,
    ) -> ElGamalCiphertext {
        ElGamal::encrypt_with(amount, self, opening)
    }

    /// Generates a decryption handle for an ElGamal public key under a Pedersen
    /// opening.
    pub fn decrypt_handle(self, opening: &PedersenOpening) -> DecryptHandle {
        DecryptHandle::new(&self, opening)
    }
}

/// Secret key for the ElGamal encryption scheme.
#[derive(Clone, Copy)]
pub struct ElGamalSecretKey(FieldValue<ScalarField>);
impl ElGamalSecretKey {
    /// Create an `ElGamalSecretKey` from a `FieldValue<ScalarField>` element.
    pub fn new(secret: FieldValue<ScalarField>) -> Self {
        assert!(!secret.is_plaintext());
        Self(secret)
    }

    /// Randomly samples an ElGamal secret key.
    ///
    /// This function is randomized. It internally samples a scalar element using a random singlet.
    pub fn new_rand() -> Self {
        ElGamalSecretKey(FieldValue::<ScalarField>::random())
    }

    pub fn get_scalar(&self) -> &FieldValue<ScalarField> {
        &self.0
    }

    /// Decrypts a ciphertext using the ElGamal secret key.
    pub fn decrypt(&self, ciphertext: &ElGamalCiphertext) -> CurveValue {
        ElGamal::decrypt(self, ciphertext)
    }
}

/// Ciphertext for the ElGamal encryption scheme.
#[derive(Clone, Copy)]
pub struct ElGamalCiphertext {
    pub commitment: PedersenCommitment,
    pub handle: DecryptHandle,
}
impl ElGamalCiphertext {
    /// Decrypts the ciphertext using an ElGamal secret key.
    pub fn decrypt(&self, secret: &ElGamalSecretKey) -> CurveValue {
        ElGamal::decrypt(secret, self)
    }
}

impl Add for ElGamalCiphertext {
    type Output = ElGamalCiphertext;

    fn add(self, ciphertext: ElGamalCiphertext) -> Self::Output {
        ElGamalCiphertext {
            commitment: self.commitment + ciphertext.commitment,
            handle: self.handle + ciphertext.handle,
        }
    }
}

impl Sub for ElGamalCiphertext {
    type Output = ElGamalCiphertext;

    fn sub(self, ciphertext: ElGamalCiphertext) -> Self::Output {
        ElGamalCiphertext {
            commitment: self.commitment - ciphertext.commitment,
            handle: self.handle - ciphertext.handle,
        }
    }
}

impl Mul<ElGamalCiphertext> for FieldValue<ScalarField> {
    type Output = ElGamalCiphertext;

    fn mul(self, rhs: ElGamalCiphertext) -> Self::Output {
        ElGamalCiphertext {
            commitment: self * rhs.commitment,
            handle: self * rhs.handle,
        }
    }
}

/// Decryption handle for Pedersen commitment.
#[derive(Clone, Copy)]
pub struct DecryptHandle(CurveValue);
impl DecryptHandle {
    pub fn new(public: &ElGamalPubkey, opening: &PedersenOpening) -> Self {
        Self((*opening.get_scalar() * public.0).reveal())
    }

    /// Create a `DecryptHandle` from a `CurveValue`.
    pub fn new_from_inner(handle: CurveValue) -> Self {
        assert!(handle.is_plaintext());
        Self(handle)
    }

    pub fn get_point(&self) -> &CurveValue {
        &self.0
    }
}

impl Add for DecryptHandle {
    type Output = DecryptHandle;

    fn add(self, handle: DecryptHandle) -> Self::Output {
        DecryptHandle(self.0 + handle.0)
    }
}

impl Sub for DecryptHandle {
    type Output = DecryptHandle;

    fn sub(self, handle: DecryptHandle) -> Self::Output {
        DecryptHandle(self.0 - handle.0)
    }
}

impl Mul<DecryptHandle> for FieldValue<ScalarField> {
    type Output = DecryptHandle;

    fn mul(self, rhs: DecryptHandle) -> Self::Output {
        DecryptHandle(self * rhs.0)
    }
}