[zk-token-sdk] clean up pod ciphertext arithmetic operations (#28470)

* clean up pod ciphertext arithmetic operations

* rename add_ciphertext to add

zk-token-sdk/src/zk_token_elgamal/ops.rs

@@ -1,116 +1,23 @@
-pub use target_arch::*;
+use crate::{
-
-#[cfg(not(target_os = "solana"))]
-mod target_arch {
-    use {
-        crate::{encryption::elgamal::ElGamalCiphertext, zk_token_elgamal::pod},
-        curve25519_dalek::{constants::RISTRETTO_BASEPOINT_COMPRESSED, scalar::Scalar},
-        std::convert::TryInto,
-    };
-    pub const TWO_32: u64 = 4294967296;
-
-    // On input two scalars x0, x1 and two ciphertexts ct0, ct1,
-    // returns `Some(x0*ct0 + x1*ct1)` or `None` if the input was invalid
-    fn add_ciphertexts(
-        scalar_0: Scalar,
-        ct_0: &pod::ElGamalCiphertext,
-        scalar_1: Scalar,
-        ct_1: &pod::ElGamalCiphertext,
-    ) -> Option<pod::ElGamalCiphertext> {
-        let ct_0: ElGamalCiphertext = (*ct_0).try_into().ok()?;
-        let ct_1: ElGamalCiphertext = (*ct_1).try_into().ok()?;
-
-        let ct_sum = ct_0 * scalar_0 + ct_1 * scalar_1;
-        Some(pod::ElGamalCiphertext::from(ct_sum))
-    }
-
-    pub(crate) fn combine_lo_hi(
-        ct_lo: &pod::ElGamalCiphertext,
-        ct_hi: &pod::ElGamalCiphertext,
-    ) -> Option<pod::ElGamalCiphertext> {
-        add_ciphertexts(Scalar::one(), ct_lo, Scalar::from(TWO_32), ct_hi)
-    }
-
-    pub fn add(
-        ct_0: &pod::ElGamalCiphertext,
-        ct_1: &pod::ElGamalCiphertext,
-    ) -> Option<pod::ElGamalCiphertext> {
-        add_ciphertexts(Scalar::one(), ct_0, Scalar::one(), ct_1)
-    }
-
-    pub fn add_with_lo_hi(
-        ct_0: &pod::ElGamalCiphertext,
-        ct_1_lo: &pod::ElGamalCiphertext,
-        ct_1_hi: &pod::ElGamalCiphertext,
-    ) -> Option<pod::ElGamalCiphertext> {
-        let ct_1 = combine_lo_hi(ct_1_lo, ct_1_hi)?;
-        add_ciphertexts(Scalar::one(), ct_0, Scalar::one(), &ct_1)
-    }
-
-    pub fn subtract(
-        ct_0: &pod::ElGamalCiphertext,
-        ct_1: &pod::ElGamalCiphertext,
-    ) -> Option<pod::ElGamalCiphertext> {
-        add_ciphertexts(Scalar::one(), ct_0, -Scalar::one(), ct_1)
-    }
-
-    pub fn subtract_with_lo_hi(
-        ct_0: &pod::ElGamalCiphertext,
-        ct_1_lo: &pod::ElGamalCiphertext,
-        ct_1_hi: &pod::ElGamalCiphertext,
-    ) -> Option<pod::ElGamalCiphertext> {
-        let ct_1 = combine_lo_hi(ct_1_lo, ct_1_hi)?;
-        add_ciphertexts(Scalar::one(), ct_0, -Scalar::one(), &ct_1)
-    }
-
-    pub fn add_to(ct: &pod::ElGamalCiphertext, amount: u64) -> Option<pod::ElGamalCiphertext> {
-        let mut amount_as_ct = [0_u8; 64];
-        amount_as_ct[..32].copy_from_slice(RISTRETTO_BASEPOINT_COMPRESSED.as_bytes());
-        add_ciphertexts(
-            Scalar::one(),
-            ct,
-            Scalar::from(amount),
-            &pod::ElGamalCiphertext(amount_as_ct),
-        )
-    }
-
-    pub fn subtract_from(
-        ct: &pod::ElGamalCiphertext,
-        amount: u64,
-    ) -> Option<pod::ElGamalCiphertext> {
-        let mut amount_as_ct = [0_u8; 64];
-        amount_as_ct[..32].copy_from_slice(RISTRETTO_BASEPOINT_COMPRESSED.as_bytes());
-        add_ciphertexts(
-            Scalar::one(),
-            ct,
-            -Scalar::from(amount),
-            &pod::ElGamalCiphertext(amount_as_ct),
-        )
-    }
-}
-
-#[cfg(target_os = "solana")]
-#[allow(unused_variables)]
-mod target_arch {
-    use crate::{
     curve25519::{
         ristretto::{add_ristretto, multiply_ristretto, subtract_ristretto, PodRistrettoPoint},
         scalar::PodScalar,
     },
     zk_token_elgamal::pod,
-    };
+};
 
-    const SHIFT_BITS: usize = 16;
+const SHIFT_BITS: usize = 16;
 
-    const G: PodRistrettoPoint = PodRistrettoPoint([
+const G: PodRistrettoPoint = PodRistrettoPoint([
-        226, 242, 174, 10, 106, 188, 78, 113, 168, 132, 169, 97, 197, 0, 81, 95, 88, 227, 11, 106,
+    226, 242, 174, 10, 106, 188, 78, 113, 168, 132, 169, 97, 197, 0, 81, 95, 88, 227, 11, 106, 165,
-        165, 130, 221, 141, 182, 166, 89, 69, 224, 141, 45, 118,
+    130, 221, 141, 182, 166, 89, 69, 224, 141, 45, 118,
-    ]);
+]);
 
-    pub fn add(
+/// Add two ElGamal ciphertexts
+pub fn add(
     left_ciphertext: &pod::ElGamalCiphertext,
     right_ciphertext: &pod::ElGamalCiphertext,
-    ) -> Option<pod::ElGamalCiphertext> {
+) -> Option<pod::ElGamalCiphertext> {
     let (left_commitment, left_handle): (pod::PedersenCommitment, pod::DecryptHandle) =
         (*left_ciphertext).into();
     let (right_commitment, right_handle): (pod::PedersenCommitment, pod::DecryptHandle) =
@@ -122,23 +29,42 @@ mod target_arch {
         add_ristretto(&left_handle.into(), &right_handle.into())?.into();
 
     Some((result_commitment, result_handle).into())
-    }
+}
 
-    pub fn add_with_lo_hi(
+/// Multiply an ElGamal ciphertext by a scalar
+pub fn multiply(
+    scalar: &PodScalar,
+    ciphertext: &pod::ElGamalCiphertext,
+) -> Option<pod::ElGamalCiphertext> {
+    let (commitment, handle): (pod::PedersenCommitment, pod::DecryptHandle) = (*ciphertext).into();
+
+    let commitment_point: PodRistrettoPoint = commitment.into();
+    let handle_point: PodRistrettoPoint = handle.into();
+
+    let result_commitment: pod::PedersenCommitment =
+        multiply_ristretto(scalar, &commitment_point)?.into();
+    let result_handle: pod::DecryptHandle = multiply_ristretto(scalar, &handle_point)?.into();
+
+    Some((result_commitment, result_handle).into())
+}
+
+/// Compute `left_ciphertext + (right_ciphertext_lo + 2^16 * right_ciphertext_hi)`
+pub fn add_with_lo_hi(
     left_ciphertext: &pod::ElGamalCiphertext,
     right_ciphertext_lo: &pod::ElGamalCiphertext,
     right_ciphertext_hi: &pod::ElGamalCiphertext,
-    ) -> Option<pod::ElGamalCiphertext> {
+) -> Option<pod::ElGamalCiphertext> {
     let shift_scalar = to_scalar(1_u64 << SHIFT_BITS);
-        let shifted_right_ciphertext_hi = scalar_ciphertext(&shift_scalar, &right_ciphertext_hi)?;
+    let shifted_right_ciphertext_hi = multiply(&shift_scalar, right_ciphertext_hi)?;
     let combined_right_ciphertext = add(right_ciphertext_lo, &shifted_right_ciphertext_hi)?;
     add(left_ciphertext, &combined_right_ciphertext)
-    }
+}
 
-    pub fn subtract(
+/// Subtract two ElGamal ciphertexts
+pub fn subtract(
     left_ciphertext: &pod::ElGamalCiphertext,
     right_ciphertext: &pod::ElGamalCiphertext,
-    ) -> Option<pod::ElGamalCiphertext> {
+) -> Option<pod::ElGamalCiphertext> {
     let (left_commitment, left_handle): (pod::PedersenCommitment, pod::DecryptHandle) =
         (*left_ciphertext).into();
     let (right_commitment, right_handle): (pod::PedersenCommitment, pod::DecryptHandle) =
@@ -150,78 +76,56 @@ mod target_arch {
         subtract_ristretto(&left_handle.into(), &right_handle.into())?.into();
 
     Some((result_commitment, result_handle).into())
-    }
+}
 
-    pub fn subtract_with_lo_hi(
+/// Compute `left_ciphertext - (right_ciphertext_lo + 2^16 * right_ciphertext_hi)`
+pub fn subtract_with_lo_hi(
     left_ciphertext: &pod::ElGamalCiphertext,
     right_ciphertext_lo: &pod::ElGamalCiphertext,
     right_ciphertext_hi: &pod::ElGamalCiphertext,
-    ) -> Option<pod::ElGamalCiphertext> {
+) -> Option<pod::ElGamalCiphertext> {
     let shift_scalar = to_scalar(1_u64 << SHIFT_BITS);
-        let shifted_right_ciphertext_hi = scalar_ciphertext(&shift_scalar, &right_ciphertext_hi)?;
+    let shifted_right_ciphertext_hi = multiply(&shift_scalar, right_ciphertext_hi)?;
     let combined_right_ciphertext = add(right_ciphertext_lo, &shifted_right_ciphertext_hi)?;
     subtract(left_ciphertext, &combined_right_ciphertext)
-    }
+}
 
-    pub fn add_to(
+/// Add a constant amount to a ciphertext
-        ciphertext: &pod::ElGamalCiphertext,
+pub fn add_to(ciphertext: &pod::ElGamalCiphertext, amount: u64) -> Option<pod::ElGamalCiphertext> {
-        amount: u64,
-    ) -> Option<pod::ElGamalCiphertext> {
     let amount_scalar = to_scalar(amount);
     let amount_point = multiply_ristretto(&amount_scalar, &G)?;
 
-        let (commitment, handle): (pod::PedersenCommitment, pod::DecryptHandle) =
+    let (commitment, handle): (pod::PedersenCommitment, pod::DecryptHandle) = (*ciphertext).into();
-            (*ciphertext).into();
     let commitment_point: PodRistrettoPoint = commitment.into();
 
     let result_commitment: pod::PedersenCommitment =
         add_ristretto(&commitment_point, &amount_point)?.into();
     Some((result_commitment, handle).into())
-    }
+}
 
-    pub fn subtract_from(
+/// Subtract a constant amount to a ciphertext
+pub fn subtract_from(
     ciphertext: &pod::ElGamalCiphertext,
     amount: u64,
-    ) -> Option<pod::ElGamalCiphertext> {
+) -> Option<pod::ElGamalCiphertext> {
     let amount_scalar = to_scalar(amount);
     let amount_point = multiply_ristretto(&amount_scalar, &G)?;
 
-        let (commitment, handle): (pod::PedersenCommitment, pod::DecryptHandle) =
+    let (commitment, handle): (pod::PedersenCommitment, pod::DecryptHandle) = (*ciphertext).into();
-            (*ciphertext).into();
     let commitment_point: PodRistrettoPoint = commitment.into();
 
     let result_commitment: pod::PedersenCommitment =
         subtract_ristretto(&commitment_point, &amount_point)?.into();
     Some((result_commitment, handle).into())
-    }
+}
 
-    fn to_scalar(amount: u64) -> PodScalar {
+/// Convert a `u64` amount into a curve25519 scalar
+fn to_scalar(amount: u64) -> PodScalar {
     let mut bytes = [0u8; 32];
     bytes[..8].copy_from_slice(&amount.to_le_bytes());
     PodScalar(bytes)
-    }
-
-    fn scalar_ciphertext(
-        scalar: &PodScalar,
-        ciphertext: &pod::ElGamalCiphertext,
-    ) -> Option<pod::ElGamalCiphertext> {
-        let (commitment, handle): (pod::PedersenCommitment, pod::DecryptHandle) =
-            (*ciphertext).into();
-
-        let commitment_point: PodRistrettoPoint = commitment.into();
-        let handle_point: PodRistrettoPoint = handle.into();
-
-        let result_commitment: pod::PedersenCommitment =
-            multiply_ristretto(scalar, &commitment_point)?.into();
-        let result_handle: pod::DecryptHandle = multiply_ristretto(scalar, &handle_point)?.into();
-
-        Some((result_commitment, result_handle).into())
-    }
 }
 
-pub const OP_ADD: u64 = 0;
-pub const OP_SUB: u64 = 1;
-
 #[cfg(test)]
 mod tests {
     use {
@@ -230,6 +134,7 @@ mod tests {
                 elgamal::{ElGamalCiphertext, ElGamalKeypair},
                 pedersen::{Pedersen, PedersenOpening},
             },
+            instruction::split_u64,
             zk_token_elgamal::{ops, pod},
         },
         bytemuck::Zeroable,
@@ -237,6 +142,8 @@ mod tests {
         std::convert::TryInto,
     };
 
+    const TWO_16: u64 = 65536;
+
     #[test]
     fn test_zero_ct() {
         let spendable_balance = pod::ElGamalCiphertext::zeroed();
@@ -290,19 +197,11 @@ mod tests {
         assert_eq!(expected, subtracted_ct);
     }
 
-    /// Split u64 number into two u32 numbers
-    fn split_u64_into_u32(amt: u64) -> (u32, u32) {
-        let lo = amt as u32;
-        let hi = (amt >> 32) as u32;
-
-        (lo, hi)
-    }
-
     #[test]
     fn test_transfer_arithmetic() {
         // transfer amount
         let transfer_amount: u64 = 55;
-        let (amount_lo, amount_hi) = split_u64_into_u32(transfer_amount);
+        let (amount_lo, amount_hi) = split_u64(transfer_amount, 16, 32).unwrap();
 
         // generate public keys
         let source_pk = ElGamalKeypair::new_rand().public;
@@ -335,38 +234,26 @@ mod tests {
             dest_pk.encrypt_with(77_u64, &dest_open).into();
 
         // program arithmetic for the source account
-
-        // 1. Combine commitments and handles
         let source_lo_ct: pod::ElGamalCiphertext = (comm_lo, handle_source_lo).into();
         let source_hi_ct: pod::ElGamalCiphertext = (comm_hi, handle_source_hi).into();
 
-        // 2. Combine lo and hi ciphertexts
-        let source_combined_ct = ops::combine_lo_hi(&source_lo_ct, &source_hi_ct).unwrap();
-
-        // 3. Subtract from available balance
         let final_source_spendable =
-            ops::subtract(&source_spendable_ct, &source_combined_ct).unwrap();
+            ops::subtract_with_lo_hi(&source_spendable_ct, &source_lo_ct, &source_hi_ct).unwrap();
 
-        // test
         let final_source_open =
-            source_open - (open_lo.clone() + open_hi.clone() * Scalar::from(ops::TWO_32));
+            source_open - (open_lo.clone() + open_hi.clone() * Scalar::from(TWO_16));
         let expected_source: pod::ElGamalCiphertext =
             source_pk.encrypt_with(22_u64, &final_source_open).into();
         assert_eq!(expected_source, final_source_spendable);
 
-        // same for the destination account
+        // program arithemtic for the destination account
-
-        // 1. Combine commitments and handles
         let dest_lo_ct: pod::ElGamalCiphertext = (comm_lo, handle_dest_lo).into();
         let dest_hi_ct: pod::ElGamalCiphertext = (comm_hi, handle_dest_hi).into();
 
-        // 2. Combine lo and hi ciphertexts
+        let final_dest_pending =
-        let dest_combined_ct = ops::combine_lo_hi(&dest_lo_ct, &dest_hi_ct).unwrap();
+            ops::add_with_lo_hi(&dest_pending_ct, &dest_lo_ct, &dest_hi_ct).unwrap();
 
-        // 3. Add to pending balance
+        let final_dest_open = dest_open + (open_lo + open_hi * Scalar::from(TWO_16));
-        let final_dest_pending = ops::add(&dest_pending_ct, &dest_combined_ct).unwrap();
-
-        let final_dest_open = dest_open + (open_lo + open_hi * Scalar::from(ops::TWO_32));
         let expected_dest_ct: pod::ElGamalCiphertext =
             dest_pk.encrypt_with(132_u64, &final_dest_open).into();
         assert_eq!(expected_dest_ct, final_dest_pending);