Rweber/zkp (#205)

Can create pure R1CS BFV proof.

Cargo.lock

@@ -147,8 +147,11 @@ version = "0.1.0"
 dependencies = [
  "ark-ff",
  "ark-poly",
+ "crypto-bigint",
  "rand",
  "rand_distr",
+ "sunscreen",
+ "sunscreen_zkp_backend",
 ]
 
 [[package]]

benchmarks/bfv_zkp/Cargo.toml

@@ -8,5 +8,8 @@ edition = "2021"
 [dependencies]
 ark-ff = "0.4.0-alpha.4"
 ark-poly = "0.4.0-alpha.4"
+crypto-bigint = "0.4.9"
 rand = "0.8.5"
-rand_distr = "0.4.3"
+rand_distr = "0.4.3"
+sunscreen = { path = "../../sunscreen" }
+sunscreen_zkp_backend = { path = "../../sunscreen_zkp_backend" }

benchmarks/bfv_zkp/src/bfv.rs

@@ -1,12 +1,18 @@
 // TODO: Remove
 #![allow(unused)]
 
-use ark_ff::{BigInt, BigInteger, Fp, FpConfig, MontBackend, MontConfig};
+use ark_ff::{BigInt, BigInteger, Field, Fp, FpConfig, MontBackend, MontConfig, PrimeField};
 use ark_poly::univariate::DensePolynomial;
+use sunscreen::{
+    types::zkp::{Mod, NativeField, RnsRingPolynomial, Scale, ToBinary, ToResidues},
+    zkp_program, Application, BackendField, Compiler, Runtime, ZkpApplication, ZkpBackend,
+    ZkpProgramInput, ZkpRuntime,
+};
+use sunscreen_zkp_backend::{bulletproofs::BulletproofsBackend, BigInt as ZkpBigInt, Proof};
 
 use crate::poly_ring::PolyRing;
 
-const POLY_DEGREE: usize = 1024;
+const POLY_DEGREE: usize = 2;
 
 #[derive(MontConfig)]
 #[modulus = "132120577"]
@@ -22,6 +28,7 @@ pub type Noise = (Poly, Poly);
 
 const CIPHER_MODULUS: u64 = 132120577;
 const PLAIN_MODULUS: u64 = 1024;
+const LOG_PLAIN_MODULUS: usize = 10;
 
 pub fn gen_keys() -> (PublicKey, PrivateKey) {
     let s = PolyRing::rand_binary(POLY_DEGREE);
@@ -188,6 +195,164 @@ fn div_round_bigint<const N: usize>(a: BigInt<N>, b: BigInt<N>) -> BigInt<N> {
     div
 }
 
+type BfvPoly<F> = RnsRingPolynomial<F, POLY_DEGREE, 1>;
+
+#[zkp_program(backend = "bulletproofs")]
+fn prove_enc<F: BackendField>(
+    m: BfvPoly<F>,
+    e_1: BfvPoly<F>,
+    e_2: BfvPoly<F>,
+    u: BfvPoly<F>,
+    #[constant] expected_c_0: BfvPoly<F>,
+    #[constant] expected_c_1: BfvPoly<F>,
+    #[constant] p_0: BfvPoly<F>,
+    #[constant] p_1: BfvPoly<F>,
+    #[constant] delta: NativeField<F>,
+) {
+    let q = NativeField::<F>::from(CIPHER_MODULUS).into_program_node();
+
+    fn log2(x: usize) -> usize {
+        let log2 = 8 * std::mem::size_of::<usize>() - x.leading_zeros() as usize;
+
+        if x.is_power_of_two() {
+            log2
+        } else {
+            log2 + 1
+        }
+    }
+
+    let log_q = log2(CIPHER_MODULUS as usize);
+
+    let c_0 = m.clone().scale(delta) + p_0 * u.clone() + e_1.clone();
+    let c_0 = RnsRingPolynomial::signed_reduce(c_0, q, log_q);
+
+    let c_1 = p_1 * u.clone() + e_2.clone();
+    let c_1 = RnsRingPolynomial::signed_reduce(c_1, q, log_q);
+
+    // e_* coefficients are gaussian distributed from -19 to 19.
+    // If we add 18 to these values, we get a distribution from
+    // [0, 36], which we can range check.
+    let chi_offset = NativeField::from(19).into_program_node();
+
+    for i in 0..1 {
+        for j in 0..POLY_DEGREE {
+            // Check that u_* in [0, 2), m_* in [0, P),
+            // e_1_* and e_2_* in [0, 32)
+            u.residues()[i][j].to_unsigned::<1>();
+            m.residues()[i][j].to_unsigned::<LOG_PLAIN_MODULUS>();
+            (e_1.residues()[i][j] + chi_offset).to_unsigned::<27>();
+            (e_2.residues()[i][j] + chi_offset).to_unsigned::<27>();
+
+            c_0.residues()[i][j].constrain_eq(expected_c_0.residues()[i][j]);
+            c_1.residues()[i][j].constrain_eq(expected_c_1.residues()[i][j]);
+        }
+    }
+}
+
+pub fn compile_proof() -> ZkpApplication {
+    Compiler::new()
+        .zkp_backend::<BulletproofsBackend>()
+        .zkp_program(prove_enc)
+        .compile()
+        .unwrap()
+}
+
+fn ark_bigint_to_native_field<B: ZkpBackend, F: MontConfig<N>, const N: usize>(
+    x: Fp<MontBackend<F, N>, N>,
+) -> NativeField<B::Field> {
+    assert!(N <= 8);
+
+    let x = x.into_bigint();
+    let mut out = ZkpBigInt::ZERO;
+
+    for (i, l) in x.0.iter().enumerate() {
+        out.0.limbs_mut()[i] = crypto_bigint::Limb(*l);
+    }
+
+    NativeField::from(out)
+}
+
+type BpBackendField = <BulletproofsBackend as ZkpBackend>::Field;
+type BpField = NativeField<BpBackendField>;
+
+fn public_bfv_proof_params(
+    ciphertext: &Ciphertext,
+    public_key: &PublicKey,
+) -> Vec<ZkpProgramInput> {
+    let (p_0, p_1) = public_key.clone();
+    let (c_0, c_1) = ciphertext.clone();
+
+    let p_0 = into_rns_poly(p_0);
+    let p_1 = into_rns_poly(p_1);
+
+    let c_0 = into_rns_poly(c_0);
+    let c_1 = into_rns_poly(c_1);
+
+    let delta = BpField::from(CIPHER_MODULUS / PLAIN_MODULUS);
+
+    vec![c_0.into(), c_1.into(), p_0.into(), p_1.into(), delta.into()]
+}
+
+fn into_rns_poly<F: MontConfig<N>, const N: usize>(
+    x: PolyRing<F, N>,
+) -> RnsRingPolynomial<BpBackendField, POLY_DEGREE, 1> {
+    let mut coeffs = x
+        .poly
+        .iter()
+        .map(|x| ark_bigint_to_native_field::<BulletproofsBackend, _, N>(*x))
+        .collect::<Vec<BpField>>();
+
+    coeffs.resize(POLY_DEGREE, BpField::from(0u8));
+
+    let coeffs: [BpField; POLY_DEGREE] = coeffs.try_into().unwrap();
+
+    RnsRingPolynomial::from([coeffs])
+}
+
+pub fn prove_public_encryption(
+    app: &ZkpApplication,
+    message: &Poly,
+    encryption_data: &(Ciphertext, Noise, Poly),
+    public_key: &PublicKey,
+) -> Proof {
+    let runtime = Runtime::new_zkp(&BulletproofsBackend::new()).unwrap();
+
+    let prog = app.get_zkp_program(prove_enc).unwrap();
+
+    let (ciphertext, (e_1, e_2), u) = encryption_data.clone();
+
+    let (p_0, p_1) = public_key.clone();
+
+    let m = into_rns_poly(message.clone());
+
+    let e_1 = into_rns_poly(e_1);
+    let e_2 = into_rns_poly(e_2);
+    let u = into_rns_poly(u);
+
+    let private_args: Vec<ZkpProgramInput> = vec![m.into(), e_1.into(), e_2.into(), u.into()];
+
+    let const_args = public_bfv_proof_params(&ciphertext, public_key);
+
+    runtime
+        .prove(prog, const_args, vec![], private_args)
+        .unwrap()
+}
+
+pub fn verify_public_encryption(
+    app: &ZkpApplication,
+    proof: &Proof,
+    ciphertext: &Ciphertext,
+    public_key: &PublicKey,
+) {
+    let runtime = Runtime::new_zkp(&BulletproofsBackend::new()).unwrap();
+
+    let const_args = public_bfv_proof_params(ciphertext, public_key);
+
+    let program = app.get_zkp_program(prove_enc).unwrap();
+
+    runtime.verify(program, proof, const_args, vec![]).unwrap();
+}
+
 #[test]
 fn can_bigint_mul() {
     let test_mul = |a_u32: u32, b_u32: u32| {
@@ -270,6 +435,25 @@ fn can_encrypt_decrypt() {
     assert_eq!(m.poly, message.poly);
 }
 
+#[test]
+fn can_prove_encryption() {
+    let (public, private) = gen_keys();
+
+    let mut message = Poly::zero();
+
+    for i in 0..POLY_DEGREE {
+        message.poly.coeffs.push(Fp::from(i as u32));
+    }
+
+    let res = encrypt(&public, &message, POLY_DEGREE);
+
+    let app = compile_proof();
+
+    let proof = prove_public_encryption(&app, &message, &res, &public);
+
+    verify_public_encryption(&app, &proof, &res.0, &public);
+}
+
 #[test]
 fn can_roundtrip_ints() {
     let x = Fq::from(1234u32);

seal_fhe/src/context.rs

@@ -54,7 +54,7 @@ impl Context {
      * Creates an instance of SEALContext and performs several pre-computations
      * on the given EncryptionParameters.
      *
-     * * `params` - The encryption parameters.</param>
+     * * `params` - The encryption parameters.
      * * `expand_mod_chain` - Determines whether the modulus switching chain
      * should be created.
      * * `security_level` - Determines whether a specific security level should be

seal_fhe/src/encryption_parameters.rs

@@ -56,7 +56,7 @@ impl SchemeType {
  * data lookup and input validity checks. In modulus switching the user can use
  * the ParmsId to keep track of the chain of encryption parameters. The ParmsId is
  * not exposed in the public API of EncryptionParameters, but can be accessed
- * through the <see cref="SEALContext.ContextData" /> class once the SEALContext
+ * through the SEALContext.ContextData" class once the SEALContext
  * has been created.
  *
  * Choosing inappropriate encryption parameters may lead to an encryption scheme

seal_fhe/src/encryptor_decryptor.rs

@@ -148,9 +148,9 @@ impl Decryptor {
     /**
      * Creates a Decryptor instance initialized with the specified SEALContext
      * and secret key.
-     * </summary>
-     * <param name="context">The SEALContext</param>
-     * <param name="secretKey">The secret key</param>
+     *
+     * The SEALContext
+     * The secret key
      */
     pub fn new(ctx: &Context, secret_key: &SecretKey) -> Result<Self> {
         let mut handle = null_mut();

seal_fhe/src/key_generator.rs

@@ -24,7 +24,7 @@ unsafe impl Send for KeyGenerator {}
 impl KeyGenerator {
     /**
      *
-     * Creates a KeyGenerator initialized with the specified <see cref="SEALContext" />.
+     * Creates a KeyGenerator initialized with the specified SEALContext.
      * Dynamically allocated member variables are allocated from the global memory pool.
      *
      * * `context` - The context describing the encryption scheme.
@@ -157,13 +157,11 @@ impl KeyGenerator {
 
     /**
      * Generates Galois keys and stores the result in destination.
-     * </summary>
-     * <remarks>
-     * <para>
+     *
+     * # Remarks
      * Generates Galois keys and stores the result in destination. Every time
      * this function is called, new Galois keys will be generated.
-     * </para>
-     * <para>
+     *
      * This function creates logarithmically many (in degree of the polynomial
      * modulus) Galois keys that is sufficient to apply any Galois automorphism
      * (e.g. rotations) on encrypted data. Most users will want to use this

sunscreen/src/lib.rs

@@ -83,7 +83,7 @@ pub use sunscreen_runtime::{
     CallSignature, Ciphertext, CompiledFheProgram, Error as RuntimeError, FheProgramInput,
     FheProgramInputTrait, FheProgramMetadata, FheRuntime, FheZkpRuntime, InnerCiphertext,
     InnerPlaintext, Params, Plaintext, PrivateKey, PublicKey, RequiredKeys, Runtime, WithContext,
-    ZkpRuntime,
+    ZkpProgramInput, ZkpRuntime,
 };
 pub use sunscreen_zkp_backend::{BackendField, Error as ZkpError, Result as ZkpResult, ZkpBackend};
 pub use zkp::ZkpProgramFn;

sunscreen/src/types/zkp/gadgets/arithmetic.rs

@@ -24,9 +24,11 @@ impl SignedModulus {
      *
      * # Panics
      * * When `field_modulus == 0`
+     * * When max_remainder_bits > 512
      */
     pub fn new(field_modulus: BigInt, max_remainder_bits: usize) -> Self {
         assert_ne!(field_modulus, BigInt::ZERO);
+        assert!(max_remainder_bits <= 512);
 
         Self {
             field_modulus,

sunscreen/src/types/zkp/gadgets/binary.rs

@@ -10,7 +10,17 @@ pub struct ToUInt {
 }
 
 impl ToUInt {
+    /**
+     * Creates a new [`ToUInt`] gadget.
+     *
+     * # Panics
+     * * If n > 512
+     */
     pub fn new(n: usize) -> Self {
+        if n > 512 {
+            panic!("Cannot decompose into > 512 bit values.");
+        }
+
         Self { n }
     }
 }
@@ -23,6 +33,12 @@ impl Gadget for ToUInt {
             return Err(ZkpError::gadget_error("Cannot create 0-bit uint."));
         }
 
+        if self.n > 512 {
+            return Err(ZkpError::gadget_error(
+                "Cannot decompose into > 512-bit values.",
+            ));
+        }
+
         if *val > BigInt::ONE.shl_vartime(self.n) {
             return Err(ZkpError::gadget_error(&format!(
                 "Value too large for {} bit unsigned int.",

sunscreen/src/types/zkp/mod.rs

@@ -6,6 +6,7 @@ mod rns_polynomial;
 pub use native_field::*;
 use petgraph::stable_graph::NodeIndex;
 pub use program_node::*;
+pub use rns_polynomial::*;
 use sunscreen_compiler_common::TypeName;
 
 pub use sunscreen_runtime::{ToNativeFields, ZkpProgramInputTrait};

sunscreen/src/types/zkp/native_field.rs

@@ -52,6 +52,15 @@ impl<F: BackendField> NativeField<F> {
     }
 }
 
+impl<F: BackendField> From<BigInt> for NativeField<F> {
+    fn from(val: BigInt) -> Self {
+        Self {
+            val,
+            _phantom: PhantomData,
+        }
+    }
+}
+
 impl<F: BackendField> From<u8> for NativeField<F> {
     fn from(x: u8) -> Self {
         (u64::from(x)).into()
@@ -244,6 +253,9 @@ where
      * Additionally, this value must be less than `log2(f)` where `f`
      * is the size of the backend field.
      *
+     * # Panics
+     * Implementors should panic if remainder_bits > 512.
+     *
      * # Example
      * Suppose the native field is F_11 and the desired field is F_7
      * (i.e. m = 7).

sunscreen/src/types/zkp/program_node.rs

@@ -196,7 +196,7 @@ pub trait ConstrainEq<Rhs> {
 
 impl<T, U, V> ConstrainEq<T> for U
 where
-    T: ZkpType + Sized + IntoProgramNode<Output = V>,
+    T: Sized + IntoProgramNode<Output = V>,
     U: IntoProgramNode<Output = V> + Sized,
     V: ZkpType + Sized + ConstrainEqVarVar,
 {

sunscreen/src/types/zkp/rns_polynomial.rs

@@ -1,3 +1,4 @@
+use petgraph::stable_graph::NodeIndex;
 use sunscreen_compiler_macros::TypeName;
 use sunscreen_runtime::ZkpProgramInputTrait;
 use sunscreen_zkp_backend::{BackendField, BigInt};
@@ -8,12 +9,12 @@ use crate::{
     zkp::ZkpContextOps,
 };
 
-use super::{AddVar, MulVar, NativeField, NumFieldElements, ToNativeFields, ZkpType};
+use super::{AddVar, Mod, MulVar, NativeField, NumFieldElements, ToNativeFields, ZkpType};
 
 use crate as sunscreen;
 
 /**
- * A polynomial in Z_q[X]/(X^N+1), up to degree N-1. `q` is the
+ * A polynomial in `Z_q[X]/(X^N+1)`, up to degree N-1. `q` is the
  * coefficient modulus and is the product of R factors. Each coefficient is decomposed
  * into R residues (i.e. RNS form).
  *
@@ -56,7 +57,14 @@ impl<F: BackendField, const N: usize, const R: usize> ToNativeFields
 
 impl<F: BackendField, const N: usize, const R: usize> ZkpType for RnsRingPolynomial<F, N, R> {}
 
+/**
+ * Returns the RNS residues for each coefficient. The coefficient index
+ * is the leading dimension for efficient NTT transforms.
+ */
 pub trait ToResidues<F: BackendField, const N: usize, const R: usize> {
+    /**
+     * Return the residues.
+     */
     fn residues(&self) -> [[ProgramNode<NativeField<F>>; N]; R];
 }
 
@@ -134,15 +142,61 @@ impl<F: BackendField, const N: usize, const R: usize> MulVar for RnsRingPolynomi
     }
 }
 
+/**
+ * For scaling an algebraic structure (e.g. polynomial.)
+ */
+pub trait Scale<F: BackendField> {
+    /**
+     * Return a structure scaled by `x`.
+     */
+    fn scale(self, x: ProgramNode<NativeField<F>>) -> Self;
+}
+
+impl<F: BackendField, const D: usize, const R: usize> Scale<F>
+    for ProgramNode<RnsRingPolynomial<F, D, R>>
+{
+    fn scale(self, x: ProgramNode<NativeField<F>>) -> Self {
+        let mut output = vec![NodeIndex::from(0); D * R];
+
+        with_zkp_ctx(|ctx| {
+            for (i, o) in output.iter_mut().enumerate().take(R * D) {
+                *o = ctx.add_multiplication(self.ids[i], x.ids[0]);
+            }
+        });
+
+        Self::new(&output)
+    }
+}
+
+impl<F: BackendField, const D: usize, const R: usize> Mod<F> for RnsRingPolynomial<F, D, R> {
+    fn signed_reduce(
+        lhs: ProgramNode<Self>,
+        m: ProgramNode<NativeField<F>>,
+        remainder_bits: usize,
+    ) -> ProgramNode<Self> {
+        let residues = lhs.residues();
+
+        let mut outputs = vec![];
+
+        for r in residues.iter().take(R) {
+            for j in r {
+                outputs.push(NativeField::signed_reduce(*j, m, remainder_bits).ids[0]);
+            }
+        }
+
+        ProgramNode::new(&outputs)
+    }
+}
+
 #[cfg(test)]
 mod tests {
-    use sunscreen_runtime::Runtime;
+    use sunscreen_runtime::{Runtime, ZkpProgramInput};
     use sunscreen_zkp_backend::bulletproofs::BulletproofsBackend;
     use sunscreen_zkp_backend::{BackendField, ZkpBackend};
 
     use crate as sunscreen;
     use crate::types::zkp::rns_polynomial::{RnsRingPolynomial, ToResidues};
-    use crate::types::zkp::NativeField;
+    use crate::types::zkp::{NativeField, Scale};
     use crate::{zkp_program, Compiler};
 
     #[test]
@@ -254,4 +308,58 @@ mod tests {
 
         assert!(result.is_err());
     }
+
+    #[test]
+    fn can_scale_polynomial() {
+        #[zkp_program(backend = "bulletproofs")]
+        fn add_poly<F: BackendField>(
+            #[constant] a: RnsRingPolynomial<F, 8, 2>,
+            #[constant] b: NativeField<F>,
+        ) {
+            let c = a.scale(b);
+
+            let expected = [
+                [2u8, 4u8, 6u8, 8u8, 10u8, 12u8, 14u8, 16u8],
+                [18u8, 20u8, 22u8, 24u8, 26u8, 28u8, 30u8, 32],
+            ];
+
+            let residues = c.residues();
+
+            for i in 0..residues.len() {
+                for j in 0..residues[i].len() {
+                    residues[i][j].constrain_eq(NativeField::from(expected[i][j]));
+                }
+            }
+        }
+
+        let app = Compiler::new()
+            .zkp_backend::<BulletproofsBackend>()
+            .zkp_program(add_poly)
+            .compile()
+            .unwrap();
+
+        let runtime = Runtime::new_zkp(&BulletproofsBackend::new()).unwrap();
+
+        let program = app.get_zkp_program(add_poly).unwrap();
+
+        type BPRnsRingPolynomial<const N: usize, const R: usize> =
+            RnsRingPolynomial<<BulletproofsBackend as ZkpBackend>::Field, N, R>;
+
+        let a = BPRnsRingPolynomial::from([
+            [1u8, 2, 3, 4, 5, 6, 7, 8],
+            [9, 10, 11, 12, 13, 14, 15, 16],
+        ]);
+
+        type BpField = NativeField<<BulletproofsBackend as ZkpBackend>::Field>;
+
+        let b = BpField::from(2u8);
+
+        let const_args: Vec<ZkpProgramInput> = vec![a.into(), b.into()];
+
+        let proof = runtime
+            .prove(program, const_args.clone(), vec![], vec![])
+            .unwrap();
+
+        runtime.verify(program, &proof, const_args, vec![]).unwrap();
+    }
 }

sunscreen_compiler_common/src/macros/mod.rs

@@ -23,7 +23,7 @@ pub enum ProgramTypeError {
 
 /**
  * Given an input type T, returns
- * * ProgramNode<T> when T is a Path
+ * * `ProgramNode<T>` when T is a Path
  * * [map_input_type(T); N] when T is Array
  */
 pub fn lift_type(arg_type: &Type) -> Result<Type, ProgramTypeError> {

sunscreen_compiler_macros/src/zkp_program.rs

@@ -159,7 +159,7 @@ fn parse_inner(_attr_params: ZkpProgramAttrs, input_fn: ItemFn) -> Result<TokenS
             fn build(&self) -> sunscreen::Result<sunscreen::ZkpFrontendCompilation> {
                 use std::cell::RefCell;
                 use std::mem::transmute;
-                use sunscreen::{CURRENT_ZKP_CTX, ZkpContext, ZkpData, Error, INDEX_ARENA, Result, types::{zkp::{ProgramNode, ConstrainEq}, TypeName}};
+                use sunscreen::{CURRENT_ZKP_CTX, ZkpContext, ZkpData, Error, INDEX_ARENA, Result, types::{zkp::{ProgramNode, ConstrainEq, IntoProgramNode}, TypeName}};
 
                 let mut context = ZkpContext::new(ZkpData::new());
 

sunscreen_runtime/src/lib.rs

@@ -12,6 +12,8 @@ mod run;
 mod runtime;
 mod serialization;
 
+use std::sync::Arc;
+
 pub use crate::error::*;
 pub use crate::keys::*;
 pub use crate::metadata::*;
@@ -203,17 +205,18 @@ pub enum FheProgramInput {
  */
 pub trait ZkpProgramInputTrait: ToNativeFields + TypeNameInstance {}
 
+#[derive(Clone)]
 /**
  * An input argument to a ZKP program.
  */
-pub struct ZkpProgramInput(pub Box<dyn ZkpProgramInputTrait>);
+pub struct ZkpProgramInput(pub Arc<dyn ZkpProgramInputTrait>);
 
 impl<T> From<T> for ZkpProgramInput
 where
     T: ZkpProgramInputTrait + 'static,
 {
     fn from(val: T) -> Self {
-        Self(Box::new(val))
+        Self(Arc::new(val))
     }
 }
 

sunscreen_zkp_backend/src/lib.rs

@@ -157,7 +157,12 @@ pub enum Proof {
  * A large integer representing a backend-agnostic
  * field element.
  */
-pub struct BigInt(U512);
+pub struct BigInt(
+    /**
+     * The wrapped value.
+     */
+    pub U512,
+);
 
 impl<T> std::convert::From<T> for BigInt
 where