arcis-compiler 0.9.7

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

use crate::{
    core::{
        circuits::boolean::{boolean_value::BooleanValue, byte::Byte},
        global_value::{
            curve_value::{CompressedCurveValue, CurveValue},
            value::FieldValue,
        },
    },
    traits::{Reveal, ToLeBytes},
    utils::{
        field::ScalarField,
        zkp::{
            elgamal::ElGamalPubkey,
            grouped_elgamal::{
                GroupedElGamalCiphertext3Handles,
                GROUPED_ELGAMAL_CIPHERTEXT_3_HANDLES_LEN,
            },
            pedersen::PedersenOpening,
            transcript::Transcript,
            util::UNIT_LEN,
        },
    },
};
use std::sync::LazyLock;
use zk_elgamal_proof::encryption::{pedersen::H, ELGAMAL_PUBKEY_LEN};

pub const BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_LEN: usize = 192;

pub const BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_CONTEXT_LEN: usize =
    3 * ELGAMAL_PUBKEY_LEN + 2 * GROUPED_ELGAMAL_CIPHERTEXT_3_HANDLES_LEN;

pub const BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_DATA_LEN: usize =
    BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_CONTEXT_LEN
        + BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_LEN;

/// The instruction data that is needed for the
/// `ProofInstruction::VerifyBatchedGroupedCiphertext3HandlesValidity` instruction.
///
/// It includes the cryptographic proof as well as the context data information needed to verify
/// the proof.
#[derive(Clone, Copy)]
pub struct BatchedGroupedCiphertext3HandlesValidityProofData {
    pub context: BatchedGroupedCiphertext3HandlesValidityProofContext,
    pub proof: BatchedGroupedCiphertext3HandlesValidityProof,
}

#[derive(Clone, Copy)]
pub struct BatchedGroupedCiphertext3HandlesValidityProofContext {
    pub first_pubkey: ElGamalPubkey,
    pub second_pubkey: ElGamalPubkey,
    pub third_pubkey: ElGamalPubkey,
    pub grouped_ciphertext_lo: GroupedElGamalCiphertext3Handles,
    pub grouped_ciphertext_hi: GroupedElGamalCiphertext3Handles,
}

impl BatchedGroupedCiphertext3HandlesValidityProofContext {
    pub fn new_transcript(&self) -> Transcript<BooleanValue> {
        let mut transcript =
            Transcript::new(b"batched-grouped-ciphertext-validity-3-handles-instruction");
        transcript.append_point(b"first-pubkey", &self.first_pubkey.get_point().compress());
        transcript.append_point(b"second-pubkey", &self.second_pubkey.get_point().compress());
        transcript.append_point(b"third-pubkey", &self.third_pubkey.get_point().compress());
        transcript.append_elgamal_ciphertext_3_handles(
            b"grouped-ciphertext-lo",
            &self.grouped_ciphertext_lo,
        );
        transcript.append_elgamal_ciphertext_3_handles(
            b"grouped-ciphertext-hi",
            &self.grouped_ciphertext_hi,
        );
        transcript
    }

    pub fn to_bytes(
        &self,
    ) -> [Byte<BooleanValue>; BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_CONTEXT_LEN] {
        let mut bytes = [Byte::<BooleanValue>::from(0);
            BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_CONTEXT_LEN];
        bytes[..ELGAMAL_PUBKEY_LEN]
            .copy_from_slice(&self.first_pubkey.get_point().compress().to_bytes());
        let mut offset = ELGAMAL_PUBKEY_LEN;
        bytes[offset..offset + ELGAMAL_PUBKEY_LEN]
            .copy_from_slice(&self.second_pubkey.get_point().compress().to_bytes());
        offset += ELGAMAL_PUBKEY_LEN;
        bytes[offset..offset + ELGAMAL_PUBKEY_LEN]
            .copy_from_slice(&self.third_pubkey.get_point().compress().to_bytes());
        offset += ELGAMAL_PUBKEY_LEN;
        bytes[offset..offset + GROUPED_ELGAMAL_CIPHERTEXT_3_HANDLES_LEN]
            .copy_from_slice(&self.grouped_ciphertext_lo.to_bytes());
        offset += GROUPED_ELGAMAL_CIPHERTEXT_3_HANDLES_LEN;
        bytes[offset..offset + GROUPED_ELGAMAL_CIPHERTEXT_3_HANDLES_LEN]
            .copy_from_slice(&self.grouped_ciphertext_hi.to_bytes());
        bytes
    }
}

impl BatchedGroupedCiphertext3HandlesValidityProofData {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        first_pubkey: &ElGamalPubkey,
        second_pubkey: &ElGamalPubkey,
        third_pubkey: &ElGamalPubkey,
        grouped_ciphertext_lo: &GroupedElGamalCiphertext3Handles,
        grouped_ciphertext_hi: &GroupedElGamalCiphertext3Handles,
        amount_lo: FieldValue<ScalarField>,
        amount_hi: FieldValue<ScalarField>,
        opening_lo: &PedersenOpening,
        opening_hi: &PedersenOpening,
    ) -> Self {
        let context = BatchedGroupedCiphertext3HandlesValidityProofContext {
            first_pubkey: *first_pubkey,
            second_pubkey: *second_pubkey,
            third_pubkey: *third_pubkey,
            grouped_ciphertext_lo: *grouped_ciphertext_lo,
            grouped_ciphertext_hi: *grouped_ciphertext_hi,
        };
        let mut transcript = context.new_transcript();
        let proof = BatchedGroupedCiphertext3HandlesValidityProof::new(
            first_pubkey,
            second_pubkey,
            third_pubkey,
            amount_lo,
            amount_hi,
            opening_lo,
            opening_hi,
            &mut transcript,
        );
        BatchedGroupedCiphertext3HandlesValidityProofData { context, proof }
    }

    pub fn to_bytes(
        &self,
    ) -> [Byte<BooleanValue>; BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_DATA_LEN] {
        let mut bytes = [Byte::<BooleanValue>::from(0);
            BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_DATA_LEN];
        bytes[..BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_CONTEXT_LEN]
            .copy_from_slice(&self.context.to_bytes());
        bytes[BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_CONTEXT_LEN..]
            .copy_from_slice(&self.proof.to_bytes());
        bytes
    }
}

/// The grouped ciphertext validity proof for 3 handles.
///
/// Contains all the elliptic curve and scalar components that make up the sigma protocol.
#[allow(non_snake_case, dead_code)]
#[derive(Clone, Copy)]
pub struct GroupedCiphertext3HandlesValidityProof {
    Y_0: CompressedCurveValue,
    Y_1: CompressedCurveValue,
    Y_2: CompressedCurveValue,
    Y_3: CompressedCurveValue,
    z_r: FieldValue<ScalarField>,
    z_x: FieldValue<ScalarField>,
}

#[allow(non_snake_case)]
impl GroupedCiphertext3HandlesValidityProof {
    /// Creates a grouped ciphertext with 3 handles validity proof.
    ///
    /// The function does *not* hash the public keys, commitment, or decryption handles into the
    /// transcript. For security, the caller (the main protocol) should hash these public
    /// components prior to invoking this constructor.
    ///
    /// This function is randomized. It uses random singlets internally to generate random scalars.
    ///
    /// Note that the proof constructor does not take the actual Pedersen commitment or decryption
    /// handles as input; it only takes the associated Pedersen opening instead.
    ///
    /// * `first_pubkey` - The first ElGamal public key
    /// * `second_pubkey` - The second ElGamal public key
    /// * `third_pubkey` - The third ElGamal public key
    /// * `amount` - The committed message in the commitment
    /// * `opening` - The opening associated with the Pedersen commitment
    /// * `transcript` - The transcript that does the bookkeeping for the Fiat-Shamir heuristic
    pub fn new(
        first_pubkey: &ElGamalPubkey,
        second_pubkey: &ElGamalPubkey,
        third_pubkey: &ElGamalPubkey,
        amount: FieldValue<ScalarField>,
        opening: &PedersenOpening,
        transcript: &mut Transcript<BooleanValue>,
    ) -> Self {
        transcript.grouped_ciphertext_validity_proof_domain_separator(3);

        // extract the relevant scalar and Ristretto points from the inputs
        let P_first = first_pubkey.get_point();
        let P_second = second_pubkey.get_point();
        let P_third = third_pubkey.get_point();

        let x = amount;
        let r = opening.get_scalar();

        // generate random masking factors that also serves as nonces
        let y_r = FieldValue::<ScalarField>::random();
        let y_x = FieldValue::<ScalarField>::random();

        let Y_0 = CurveValue::multiscalar_mul(
            vec![y_r, y_x],
            vec![
                CurveValue::from(*LazyLock::force(&H)),
                CurveValue::generator(),
            ],
        )
        .reveal()
        .compress();
        let Y_1 = (y_r * *P_first).reveal().compress();
        let Y_2 = (y_r * *P_second).reveal().compress();
        let Y_3 = (y_r * *P_third).reveal().compress();

        // record masking factors in transcript and get challenges
        transcript.append_point(b"Y_0", &Y_0);
        transcript.append_point(b"Y_1", &Y_1);
        transcript.append_point(b"Y_2", &Y_2);
        transcript.append_point(b"Y_3", &Y_3);

        let c = transcript.challenge_scalar(b"c");

        // compute masked message and opening
        let z_r = ((c * *r) + y_r).reveal();
        let z_x = ((c * x) + y_x).reveal();

        // compute challenge `w` for consistency with verification
        transcript.append_scalar(b"z_r", &z_r);
        transcript.append_scalar(b"z_x", &z_x);
        let _w = transcript.challenge_scalar(b"w");

        Self {
            Y_0,
            Y_1,
            Y_2,
            Y_3,
            z_r,
            z_x,
        }
    }
}

// Arcis implementation of https://github.com/solana-program/zk-elgamal-proof/blob/main/zk-sdk/src/sigma_proofs/batched_grouped_ciphertext_validity/handles_3.rs

/// Batched grouped ciphertext validity proof with 3 handles.
#[allow(non_snake_case, dead_code)]
#[derive(Clone, Copy)]
pub struct BatchedGroupedCiphertext3HandlesValidityProof(GroupedCiphertext3HandlesValidityProof);

#[allow(non_snake_case)]
impl BatchedGroupedCiphertext3HandlesValidityProof {
    /// Creates a batched grouped ciphertext validity proof.
    ///
    /// The function does *not* hash the public keys, commitment, or decryption handles into the
    /// transcript. For security, the caller (the main protocol) should hash these public
    /// components prior to invoking this constructor.
    ///
    /// The function simply batches the input openings and invokes the standard grouped ciphertext
    /// validity proof constructor.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        first_pubkey: &ElGamalPubkey,
        second_pubkey: &ElGamalPubkey,
        third_pubkey: &ElGamalPubkey,
        amount_lo: FieldValue<ScalarField>,
        amount_hi: FieldValue<ScalarField>,
        opening_lo: &PedersenOpening,
        opening_hi: &PedersenOpening,
        transcript: &mut Transcript<BooleanValue>,
    ) -> Self {
        transcript.batched_grouped_ciphertext_validity_proof_domain_separator(3);

        let t = transcript.challenge_scalar(b"t");

        let batched_message = amount_lo + amount_hi * t;
        let batched_opening = opening_lo + &(opening_hi * t);

        BatchedGroupedCiphertext3HandlesValidityProof(GroupedCiphertext3HandlesValidityProof::new(
            first_pubkey,
            second_pubkey,
            third_pubkey,
            batched_message,
            &batched_opening,
            transcript,
        ))
    }

    pub fn to_bytes(
        &self,
    ) -> [Byte<BooleanValue>; BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_LEN] {
        let mut buf = [Byte::<BooleanValue>::from(0);
            BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_LEN];
        let mut chunks = buf.chunks_mut(UNIT_LEN);
        chunks
            .next()
            .unwrap()
            .copy_from_slice(&self.0.Y_0.to_bytes());
        chunks
            .next()
            .unwrap()
            .copy_from_slice(&self.0.Y_1.to_bytes());
        chunks
            .next()
            .unwrap()
            .copy_from_slice(&self.0.Y_2.to_bytes());
        chunks
            .next()
            .unwrap()
            .copy_from_slice(&self.0.Y_3.to_bytes());
        chunks
            .next()
            .unwrap()
            .copy_from_slice(&self.0.z_r.to_le_bytes());
        chunks
            .next()
            .unwrap()
            .copy_from_slice(&self.0.z_x.to_le_bytes());
        buf
    }
}
#[cfg(test)]
mod tests {
    use crate::{
        core::{
            bounds::FieldBounds,
            expressions::{
                curve_expr::{CurveExpr, InputInfo},
                domain::Domain,
                expr::EvalValue,
                field_expr::FieldExpr,
                InputKind,
            },
            global_value::{
                curve_value::CurveValue,
                global_expr_store::with_local_expr_store_as_global,
                value::FieldValue,
            },
            ir_builder::{ExprStore, IRBuilder},
        },
        utils::{
            curve_point::CurvePoint,
            field::ScalarField,
            used_field::UsedField,
            zkp::{
                elgamal::{DecryptHandle, ElGamalPubkey},
                grouped_ciphertext_validity_proof::handles_3::BatchedGroupedCiphertext3HandlesValidityProofData,
                grouped_elgamal::{GroupedElGamalCiphertext, GroupedElGamalCiphertext3Handles},
                pedersen::{PedersenCommitment, PedersenOpening},
            },
        },
    };
    use group::GroupEncoding;
    use primitives::algebra::elliptic_curve::{Curve as AsyncMPCCurve, Curve25519Ristretto};
    use rand::{Rng, RngCore};
    use rustc_hash::FxHashMap;
    use std::rc::Rc;
    use zk_elgamal_proof::{
        encryption::{
            elgamal::ElGamalKeypair as SolanaElGamalKeypair,
            grouped_elgamal::GroupedElGamalCiphertext as SolanaGroupedElGamalCiphertext,
            pedersen::Pedersen as SolanaPedersen,
        },
        zk_elgamal_proof_program::proof_data::{
            BatchedGroupedCiphertext3HandlesValidityProofData as SolanaBatchedGroupedCiphertext3HandlesValidityProofData,
            ZkProofData,
        },
    };

    #[test]
    #[allow(non_snake_case)]
    fn test_ciphertext_commitment_equality() {
        let rng = &mut crate::utils::test_rng::get();

        // random ElGamal keypairs and amount
        let first_keypair = SolanaElGamalKeypair::new_rand();
        let first_pubkey = first_keypair.pubkey();

        let second_keyapir = SolanaElGamalKeypair::new_rand();
        let second_pubkey = second_keyapir.pubkey();

        let third_keypair = SolanaElGamalKeypair::new_rand();
        let third_pubkey = third_keypair.pubkey();

        let mut amount_lo = rng.next_u64();
        let mut amount_hi = rng.next_u64();

        let (commitment_lo, open_lo) = SolanaPedersen::new(amount_lo);
        let (commitment_hi, open_hi) = SolanaPedersen::new(amount_hi);

        // we flip a coin and generate an invalid proof if the bit is false
        let is_valid_proof = rng.gen_bool(0.5);
        if !is_valid_proof {
            amount_lo = rng.next_u64();
            amount_hi = rng.next_u64();
        }

        let first_handle_lo = first_pubkey.decrypt_handle(&open_lo);
        let first_handle_hi = first_pubkey.decrypt_handle(&open_hi);

        let second_handle_lo = second_pubkey.decrypt_handle(&open_lo);
        let second_handle_hi = second_pubkey.decrypt_handle(&open_hi);

        let third_handle_lo = third_pubkey.decrypt_handle(&open_lo);
        let third_handle_hi = third_pubkey.decrypt_handle(&open_hi);

        let solana_proof_data = SolanaBatchedGroupedCiphertext3HandlesValidityProofData::new(
            first_pubkey,
            second_pubkey,
            third_pubkey,
            &SolanaGroupedElGamalCiphertext {
                commitment: commitment_lo,
                handles: [first_handle_lo, second_handle_lo, third_handle_lo],
            },
            &SolanaGroupedElGamalCiphertext {
                commitment: commitment_hi,
                handles: [first_handle_hi, second_handle_hi, third_handle_hi],
            },
            amount_lo,
            amount_hi,
            &open_lo,
            &open_hi,
        )
        .unwrap();
        assert_eq!(solana_proof_data.verify_proof().is_ok(), is_valid_proof);

        let mut expr_store = IRBuilder::new(true);

        // add inputs
        let mut input_vals = FxHashMap::<usize, EvalValue>::default();
        [first_pubkey, second_pubkey, third_pubkey]
            .iter()
            .enumerate()
            .for_each(|(i, pubkey)| {
                let _ = <IRBuilder as ExprStore<ScalarField>>::push_curve(
                    &mut expr_store,
                    CurveExpr::Input(i, Rc::new(InputInfo::from(InputKind::Plaintext))),
                );
                input_vals.insert(
                    i,
                    EvalValue::Curve(CurvePoint::new(
                        <Curve25519Ristretto as AsyncMPCCurve>::Point::from_bytes(
                            &pubkey.get_point().to_bytes(),
                        )
                        .unwrap(),
                    )),
                );
            });
        let mut offset = 3;
        [commitment_lo, commitment_hi]
            .iter()
            .enumerate()
            .for_each(|(i, commitment)| {
                let _ = <IRBuilder as ExprStore<ScalarField>>::push_curve(
                    &mut expr_store,
                    CurveExpr::Input(i + offset, Rc::new(InputInfo::from(InputKind::Plaintext))),
                );
                input_vals.insert(
                    i + offset,
                    EvalValue::Curve(CurvePoint::new(
                        <Curve25519Ristretto as AsyncMPCCurve>::Point::from_bytes(
                            &commitment.get_point().to_bytes(),
                        )
                        .unwrap(),
                    )),
                );
            });
        offset += 2;
        [
            first_handle_lo,
            first_handle_hi,
            second_handle_lo,
            second_handle_hi,
            third_handle_lo,
            third_handle_hi,
        ]
        .iter()
        .enumerate()
        .for_each(|(i, handle)| {
            let _ = <IRBuilder as ExprStore<ScalarField>>::push_curve(
                &mut expr_store,
                CurveExpr::Input(i + offset, Rc::new(InputInfo::from(InputKind::Plaintext))),
            );
            input_vals.insert(
                i + offset,
                EvalValue::Curve(CurvePoint::new(
                    <Curve25519Ristretto as AsyncMPCCurve>::Point::from_bytes(
                        &handle.get_point().to_bytes(),
                    )
                    .unwrap(),
                )),
            );
        });
        offset += 6;
        [amount_lo, amount_hi]
            .iter()
            .enumerate()
            .for_each(|(i, amount)| {
                let _ = expr_store.push_field(FieldExpr::Input(
                    i + offset,
                    FieldBounds::new(
                        ScalarField::from(0),
                        ScalarField::power_of_two(64) - ScalarField::from(1),
                    )
                    .as_input_info(InputKind::Secret),
                ));
                input_vals.insert(i + offset, EvalValue::Scalar(ScalarField::from(*amount)));
            });
        offset += 2;
        [open_lo, open_hi]
            .iter()
            .enumerate()
            .for_each(|(i, opening)| {
                let _ = expr_store.push_field(FieldExpr::Input(
                    i + offset,
                    FieldBounds::<ScalarField>::All.as_input_info(InputKind::Secret),
                ));
                input_vals.insert(
                    i + offset,
                    EvalValue::Scalar(ScalarField::from(*opening.get_scalar())),
                );
            });

        let outputs = with_local_expr_store_as_global(
            || {
                let first_pubkey = CurveValue::new(0);
                let second_pubkey = CurveValue::new(1);
                let third_pubkey = CurveValue::new(2);
                let commitment_lo = CurveValue::new(3);
                let commitment_hi = CurveValue::new(4);
                let first_handle_lo = CurveValue::new(5);
                let first_handle_hi = CurveValue::new(6);
                let second_handle_lo = CurveValue::new(7);
                let second_handle_hi = CurveValue::new(8);
                let third_handle_lo = CurveValue::new(9);
                let third_handle_hi = CurveValue::new(10);
                let amount_lo = FieldValue::<ScalarField>::from_id(11);
                let amount_hi = FieldValue::<ScalarField>::from_id(12);
                let opening_lo = FieldValue::<ScalarField>::from_id(13);
                let opening_hi = FieldValue::<ScalarField>::from_id(14);

                let arcis_proof_data = BatchedGroupedCiphertext3HandlesValidityProofData::new(
                    &ElGamalPubkey::new_from_inner(first_pubkey),
                    &ElGamalPubkey::new_from_inner(second_pubkey),
                    &ElGamalPubkey::new_from_inner(third_pubkey),
                    &GroupedElGamalCiphertext3Handles(GroupedElGamalCiphertext {
                        commitment: PedersenCommitment::new(commitment_lo),
                        handles: [
                            DecryptHandle::new_from_inner(first_handle_lo),
                            DecryptHandle::new_from_inner(second_handle_lo),
                            DecryptHandle::new_from_inner(third_handle_lo),
                        ],
                    }),
                    &GroupedElGamalCiphertext3Handles(GroupedElGamalCiphertext {
                        commitment: PedersenCommitment::new(commitment_hi),
                        handles: [
                            DecryptHandle::new_from_inner(first_handle_hi),
                            DecryptHandle::new_from_inner(second_handle_hi),
                            DecryptHandle::new_from_inner(third_handle_hi),
                        ],
                    }),
                    amount_lo,
                    amount_hi,
                    &PedersenOpening::new(opening_lo),
                    &PedersenOpening::new(opening_hi),
                );

                arcis_proof_data
                    .to_bytes()
                    .into_iter()
                    .map(|byte| FieldValue::<ScalarField>::from(byte).get_id())
                    .collect::<Vec<usize>>()
            },
            &mut expr_store,
        );

        let ir = expr_store.into_ir(outputs);
        let result = ir
            .eval(rng, &mut input_vals)
            .map(|x| {
                x.into_iter()
                    .map(ScalarField::unwrap)
                    .collect::<Vec<ScalarField>>()
            })
            .unwrap();

        let arcis_proof_data_bytes = result
            .iter()
            .map(|byte| byte.to_le_bytes()[0])
            .collect::<Vec<u8>>();

        let arcis_proof_data = SolanaBatchedGroupedCiphertext3HandlesValidityProofData::from_bytes(
            &arcis_proof_data_bytes,
        )
        .unwrap();

        let arcis_verification = arcis_proof_data.verify_proof();

        assert_eq!(arcis_verification.is_ok(), is_valid_proof);
    }
}