[zk-sdk] Add pod modules for sigma and range proofs (#1281)

* add sigma `pod` module * clean up pod sigma proofs * add rangeproof `pod` module * clean up pod range proofs * update old constant names * update outdated comments * Update zk-sdk/src/sigma_proofs/pod.rs Co-authored-by: Jon C <[email protected]> --------- Co-authored-by: Jon C <[email protected]>
anza-xyz · May 14, 2024 · d54e808 · d54e808
1 parent c0e9149
commit d54e808
Show file tree

Hide file tree

Showing 4 changed files with 449 additions and 2 deletions.
diff --git a/zk-sdk/src/range_proof/mod.rs b/zk-sdk/src/range_proof/mod.rs
@@ -12,11 +12,11 @@
 
 #![allow(dead_code)]
 
+use crate::{RISTRETTO_POINT_LEN, SCALAR_LEN};
 #[cfg(not(target_os = "solana"))]
 use {
-    crate::encryption::pedersen::{Pedersen, PedersenCommitment, PedersenOpening},
     crate::{
-        encryption::pedersen::{G, H},
+        encryption::pedersen::{Pedersen, PedersenCommitment, PedersenOpening, G, H},
         range_proof::{
             errors::{RangeProofGenerationError, RangeProofVerificationError},
             generators::RangeProofGens,
@@ -37,13 +37,40 @@ use {
 };
 
 pub mod errors;
+pub mod pod;
+
 #[cfg(not(target_os = "solana"))]
 pub mod generators;
 #[cfg(not(target_os = "solana"))]
 pub mod inner_product;
 #[cfg(not(target_os = "solana"))]
 pub mod util;
 
+/// Byte length of a range proof excluding the inner-product proof component
+pub const RANGE_PROOF_MODULO_INNER_PRODUCT_PROOF_LEN: usize =
+    5 * RISTRETTO_POINT_LEN + 2 * SCALAR_LEN;
+
+/// Byte length of an inner-product proof for a vector of length 64
+pub const INNER_PRODUCT_PROOF_U64_LEN: usize = 448;
+
+/// Byte length of a range proof for an unsigned 64-bit number
+pub const RANGE_PROOF_U64_LEN: usize =
+    INNER_PRODUCT_PROOF_U64_LEN + RANGE_PROOF_MODULO_INNER_PRODUCT_PROOF_LEN;
+
+/// Byte length of an inner-product proof for a vector of length 128
+pub const INNER_PRODUCT_PROOF_U128_LEN: usize = 512;
+
+/// Byte length of a range proof for an unsigned 128-bit number
+pub const RANGE_PROOF_U128_LEN: usize =
+    INNER_PRODUCT_PROOF_U128_LEN + RANGE_PROOF_MODULO_INNER_PRODUCT_PROOF_LEN;
+
+/// Byte length of an inner-product proof for a vector of length 256
+pub const INNER_PRODUCT_PROOF_U256_LEN: usize = 576;
+
+/// Byte length of a range proof for an unsigned 256-bit number
+pub const RANGE_PROOF_U256_LEN: usize =
+    INNER_PRODUCT_PROOF_U256_LEN + RANGE_PROOF_MODULO_INNER_PRODUCT_PROOF_LEN;
+
 #[allow(non_snake_case)]
 #[cfg(not(target_os = "solana"))]
 #[derive(Clone)]

diff --git a/zk-sdk/src/range_proof/pod.rs b/zk-sdk/src/range_proof/pod.rs
@@ -0,0 +1,134 @@
+//! Plain Old Data types for range proofs.
+
+#[cfg(not(target_os = "solana"))]
+use crate::{
+    range_proof::{errors::RangeProofVerificationError, RangeProof},
+    UNIT_LEN,
+};
+use {
+    crate::range_proof::*,
+    bytemuck::{Pod, Zeroable},
+};
+
+/// The `RangeProof` type as a `Pod` restricted to proofs on 64-bit numbers.
+#[derive(Clone, Copy)]
+#[repr(transparent)]
+pub struct PodRangeProofU64(pub(crate) [u8; RANGE_PROOF_U64_LEN]);
+
+#[cfg(not(target_os = "solana"))]
+impl TryFrom<RangeProof> for PodRangeProofU64 {
+    type Error = RangeProofVerificationError;
+
+    fn try_from(decoded_proof: RangeProof) -> Result<Self, Self::Error> {
+        if decoded_proof.ipp_proof.serialized_size() != INNER_PRODUCT_PROOF_U64_LEN {
+            return Err(RangeProofVerificationError::Deserialization);
+        }
+
+        let mut buf = [0_u8; RANGE_PROOF_U64_LEN];
+        copy_range_proof_modulo_inner_product_proof(&decoded_proof, &mut buf);
+        buf[RANGE_PROOF_MODULO_INNER_PRODUCT_PROOF_LEN..RANGE_PROOF_U64_LEN]
+            .copy_from_slice(&decoded_proof.ipp_proof.to_bytes());
+        Ok(PodRangeProofU64(buf))
+    }
+}
+
+#[cfg(not(target_os = "solana"))]
+impl TryFrom<PodRangeProofU64> for RangeProof {
+    type Error = RangeProofVerificationError;
+
+    fn try_from(pod_proof: PodRangeProofU64) -> Result<Self, Self::Error> {
+        Self::from_bytes(&pod_proof.0)
+    }
+}
+
+/// The `RangeProof` type as a `Pod` restricted to proofs on 128-bit numbers.
+#[derive(Clone, Copy)]
+#[repr(transparent)]
+pub struct PodRangeProofU128(pub(crate) [u8; RANGE_PROOF_U128_LEN]);
+
+#[cfg(not(target_os = "solana"))]
+impl TryFrom<RangeProof> for PodRangeProofU128 {
+    type Error = RangeProofVerificationError;
+
+    fn try_from(decoded_proof: RangeProof) -> Result<Self, Self::Error> {
+        if decoded_proof.ipp_proof.serialized_size() != INNER_PRODUCT_PROOF_U128_LEN {
+            return Err(RangeProofVerificationError::Deserialization);
+        }
+
+        let mut buf = [0_u8; RANGE_PROOF_U128_LEN];
+        copy_range_proof_modulo_inner_product_proof(&decoded_proof, &mut buf);
+        buf[RANGE_PROOF_MODULO_INNER_PRODUCT_PROOF_LEN..RANGE_PROOF_U128_LEN]
+            .copy_from_slice(&decoded_proof.ipp_proof.to_bytes());
+        Ok(PodRangeProofU128(buf))
+    }
+}
+
+#[cfg(not(target_os = "solana"))]
+impl TryFrom<PodRangeProofU128> for RangeProof {
+    type Error = RangeProofVerificationError;
+
+    fn try_from(pod_proof: PodRangeProofU128) -> Result<Self, Self::Error> {
+        Self::from_bytes(&pod_proof.0)
+    }
+}
+
+/// The `RangeProof` type as a `Pod` restricted to proofs on 256-bit numbers.
+#[derive(Clone, Copy)]
+#[repr(transparent)]
+pub struct PodRangeProofU256(pub(crate) [u8; RANGE_PROOF_U256_LEN]);
+
+#[cfg(not(target_os = "solana"))]
+impl TryFrom<RangeProof> for PodRangeProofU256 {
+    type Error = RangeProofVerificationError;
+
+    fn try_from(decoded_proof: RangeProof) -> Result<Self, Self::Error> {
+        if decoded_proof.ipp_proof.serialized_size() != INNER_PRODUCT_PROOF_U256_LEN {
+            return Err(RangeProofVerificationError::Deserialization);
+        }
+
+        let mut buf = [0_u8; RANGE_PROOF_U256_LEN];
+        copy_range_proof_modulo_inner_product_proof(&decoded_proof, &mut buf);
+        buf[RANGE_PROOF_MODULO_INNER_PRODUCT_PROOF_LEN..RANGE_PROOF_U256_LEN]
+            .copy_from_slice(&decoded_proof.ipp_proof.to_bytes());
+        Ok(PodRangeProofU256(buf))
+    }
+}
+
+#[cfg(not(target_os = "solana"))]
+impl TryFrom<PodRangeProofU256> for RangeProof {
+    type Error = RangeProofVerificationError;
+
+    fn try_from(pod_proof: PodRangeProofU256) -> Result<Self, Self::Error> {
+        Self::from_bytes(&pod_proof.0)
+    }
+}
+
+#[cfg(not(target_os = "solana"))]
+fn copy_range_proof_modulo_inner_product_proof(proof: &RangeProof, buf: &mut [u8]) {
+    let mut chunks = buf.chunks_mut(UNIT_LEN);
+    chunks.next().unwrap().copy_from_slice(proof.A.as_bytes());
+    chunks.next().unwrap().copy_from_slice(proof.S.as_bytes());
+    chunks.next().unwrap().copy_from_slice(proof.T_1.as_bytes());
+    chunks.next().unwrap().copy_from_slice(proof.T_2.as_bytes());
+    chunks.next().unwrap().copy_from_slice(proof.t_x.as_bytes());
+    chunks
+        .next()
+        .unwrap()
+        .copy_from_slice(proof.t_x_blinding.as_bytes());
+    chunks
+        .next()
+        .unwrap()
+        .copy_from_slice(proof.e_blinding.as_bytes());
+}
+
+// The range proof pod types are wrappers for byte arrays, which are both `Pod` and `Zeroable`. However,
+// the marker traits `bytemuck::Pod` and `bytemuck::Zeroable` can only be derived for power-of-two
+// length byte arrays. Directly implement these traits for the range proof pod types.
+unsafe impl Zeroable for PodRangeProofU64 {}
+unsafe impl Pod for PodRangeProofU64 {}
+
+unsafe impl Zeroable for PodRangeProofU128 {}
+unsafe impl Pod for PodRangeProofU128 {}
+
+unsafe impl Zeroable for PodRangeProofU256 {}
+unsafe impl Pod for PodRangeProofU256 {}
diff --git a/zk-sdk/src/sigma_proofs/mod.rs b/zk-sdk/src/sigma_proofs/mod.rs
@@ -8,6 +8,7 @@
 #![allow(dead_code, unused_imports)]
 
 pub mod errors;
+pub mod pod;
 
 #[cfg(not(target_os = "solana"))]
 pub mod batched_grouped_ciphertext_validity;
@@ -24,6 +25,33 @@ pub mod pubkey;
 #[cfg(not(target_os = "solana"))]
 pub mod zero_ciphertext;
 
+/// Byte length of a ciphertext-commitment equality proof
+pub const CIPHERTEXT_COMMITMENT_EQUALITY_PROOF_LEN: usize = 192;
+
+/// Byte length of a ciphertext-ciphertext equality proof
+pub const CIPHERTEXT_CIPHERTEXT_EQUALITY_PROOF_LEN: usize = 224;
+
+/// Byte length of a grouped ciphertext for 2 handles validity proof
+pub const GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY_PROOF_LEN: usize = 160;
+
+/// Byte length of a grouped ciphertext for 3 handles validity proof
+pub const GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_LEN: usize = 192;
+
+/// Byte length of a batched grouped ciphertext for 2 handles validity proof
+pub const BATCHED_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY_PROOF_LEN: usize = 160;
+
+/// Byte length of a batched grouped ciphertext for 3 handles validity proof
+pub const BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_PROOF_LEN: usize = 192;
+
+/// Byte length of a zero-ciphertext proof
+pub const ZERO_CIPHERTEXT_PROOF_LEN: usize = 96;
+
+/// Byte length of a percentage with cap proof
+pub const PERCENTAGE_WITH_CAP_PROOF_LEN: usize = 256;
+
+/// Byte length of a public key validity proof
+pub const PUBKEY_VALIDITY_PROOF_LEN: usize = 64;
+
 #[cfg(not(target_os = "solana"))]
 use {
     crate::{sigma_proofs::errors::SigmaProofVerificationError, RISTRETTO_POINT_LEN, SCALAR_LEN},