Support BigUint plaintext moduli in BFV

crates/fhe/Cargo.toml

@@ -34,6 +34,7 @@ zeroize.workspace = true
 zeroize_derive.workspace = true
 ndarray.workspace = true
 thiserror.workspace = true
+num-integer = "0.1.46"
 
 [dev-dependencies]
 clap.workspace = true

crates/fhe/src/bfv/ciphertext.rs

@@ -267,7 +267,7 @@ mod tests {
             BfvParameters::default_arc(6, 16),
         ] {
             let sk = SecretKey::random(&params, &mut rng);
-            let v = params.plaintext.random_vec(params.degree(), &mut rng);
+            let v = fhe_math::zq::Modulus::new(params.plaintext()).unwrap().random_vec(params.degree(), &mut rng);
             let pt = Plaintext::try_encode(&v, Encoding::simd(), &params)?;
             let ct = sk.try_encrypt(&pt, &mut rng)?;
             let ct_proto = CiphertextProto::from(&ct);
@@ -288,7 +288,7 @@ mod tests {
             BfvParameters::default_arc(6, 16),
         ] {
             let sk = SecretKey::random(&params, &mut rng);
-            let v = params.plaintext.random_vec(params.degree(), &mut rng);
+            let v = fhe_math::zq::Modulus::new(params.plaintext()).unwrap().random_vec(params.degree(), &mut rng);
             let pt = Plaintext::try_encode(&v, Encoding::simd(), &params)?;
             let ct: Ciphertext = sk.try_encrypt(&pt, &mut rng)?;
             let ct_bytes = ct.to_bytes();
@@ -305,7 +305,7 @@ mod tests {
             BfvParameters::default_arc(6, 16),
         ] {
             let sk = SecretKey::random(&params, &mut rng);
-            let v = params.plaintext.random_vec(params.degree(), &mut rng);
+            let v = fhe_math::zq::Modulus::new(params.plaintext()).unwrap().random_vec(params.degree(), &mut rng);
             let pt = Plaintext::try_encode(&v, Encoding::simd(), &params)?;
             let ct: Ciphertext = sk.try_encrypt(&pt, &mut rng)?;
             let mut ct3 = &ct * &ct;
@@ -343,7 +343,7 @@ mod tests {
             BfvParameters::default_arc(6, 16),
         ] {
             let sk = SecretKey::random(&params, &mut rng);
-            let v = params.plaintext.random_vec(params.degree(), &mut rng);
+            let v = fhe_math::zq::Modulus::new(params.plaintext()).unwrap().random_vec(params.degree(), &mut rng);
             let pt = Plaintext::try_encode(&v, Encoding::simd(), &params)?;
             let mut ct: Ciphertext = sk.try_encrypt(&pt, &mut rng)?;
 
@@ -364,7 +364,7 @@ mod tests {
         let mut rng = rng();
         let params = BfvParameters::default_arc(2, 16);
         let sk = SecretKey::random(&params, &mut rng);
-        let v = params.plaintext.random_vec(params.degree(), &mut rng);
+        let v = fhe_math::zq::Modulus::new(params.plaintext()).unwrap().random_vec(params.degree(), &mut rng);
         let pt = Plaintext::try_encode(&v, Encoding::simd(), &params)?;
         let mut ct: Ciphertext = sk.try_encrypt(&pt, &mut rng)?;
 

crates/fhe/src/bfv/context/cipher_plain_context.rs

@@ -1,4 +1,5 @@
 use fhe_math::rq::{Context, Poly, scaler::Scaler};
+use num_bigint::BigUint;
 use std::sync::Arc;
 
 /// Stores pre-computed values relating a ciphertext and plaintext context pair.
@@ -11,10 +12,10 @@ pub struct CipherPlainContext {
     pub(crate) delta: Poly,
 
     /// Q modulo the plaintext modulus
-    pub(crate) q_mod_t: u64,
+    pub(crate) q_mod_t: BigUint,
 
     /// Threshold for centered reduction (plaintext_modulus + 1) / 2
-    pub(crate) plain_threshold: u64,
+    pub(crate) plain_threshold: BigUint,
 
     /// Scaler to map a ciphertext polynomial to the plaintext context
     pub(crate) scaler: Scaler,
@@ -33,8 +34,8 @@ impl CipherPlainContext {
         plaintext_context: &Arc<Context>,
         ciphertext_context: &Arc<Context>,
         delta: Poly,
-        q_mod_t: u64,
-        plain_threshold: u64,
+        q_mod_t: BigUint,
+        plain_threshold: BigUint,
         scaler: Scaler,
     ) -> Arc<Self> {
         Arc::new(CipherPlainContext {

crates/fhe/src/bfv/keys/evaluation_key.rs

@@ -551,9 +551,8 @@ mod tests {
                         .enable_inner_sum()?
                         .build(&mut rng)?;
 
-                        let v = params.plaintext.random_vec(params.degree(), &mut rng);
-                        let expected = params
-                            .plaintext
+                        let v = fhe_math::zq::Modulus::new(params.plaintext()).unwrap().random_vec(params.degree(), &mut rng);
+                        let expected = fhe_math::zq::Modulus::new(params.plaintext()).unwrap()
                             .reduce_u128(v.iter().map(|vi| *vi as u128).sum());
 
                         let pt = Plaintext::try_encode(
@@ -595,7 +594,7 @@ mod tests {
                         .enable_row_rotation()?
                         .build(&mut rng)?;
 
-                        let v = params.plaintext.random_vec(params.degree(), &mut rng);
+                        let v = fhe_math::zq::Modulus::new(params.plaintext()).unwrap().random_vec(params.degree(), &mut rng);
                         let row_size = params.degree() >> 1;
                         let mut expected = vec![0u64; params.degree()];
                         expected[..row_size].copy_from_slice(&v[row_size..]);
@@ -642,7 +641,7 @@ mod tests {
                             .enable_column_rotation(i)?
                             .build(&mut rng)?;
 
-                            let v = params.plaintext.random_vec(params.degree(), &mut rng);
+                            let v = fhe_math::zq::Modulus::new(params.plaintext()).unwrap().random_vec(params.degree(), &mut rng);
                             let row_size = params.degree() >> 1;
                             let mut expected = vec![0u64; params.degree()];
                             expected[..row_size - i].copy_from_slice(&v[i..row_size]);
@@ -699,7 +698,7 @@ mod tests {
 
                             assert!(ek.supports_expansion(i));
                             assert!(!ek.supports_expansion(i + 1));
-                            let v = params.plaintext.random_vec(1 << i, &mut rng);
+                            let v = fhe_math::zq::Modulus::new(params.plaintext()).unwrap().random_vec(1 << i, &mut rng);
                             let pt = Plaintext::try_encode(
                                 &v,
                                 Encoding::poly_at_level(ciphertext_level),
@@ -711,7 +710,7 @@ mod tests {
                             assert_eq!(ct2.len(), 1 << i);
                             for (vi, ct2i) in izip!(&v, &ct2) {
                                 let mut expected = vec![0u64; params.degree()];
-                                expected[0] = params.plaintext.mul(*vi, (1 << i) as u64);
+                                expected[0] = fhe_math::zq::Modulus::new(params.plaintext()).unwrap().mul(*vi, (1 << i) as u64);
                                 let pt = sk.try_decrypt(ct2i)?;
                                 assert_eq!(
                                     expected,

crates/fhe/src/bfv/keys/galois_key.rs

@@ -173,7 +173,7 @@ mod tests {
         ] {
             for _ in 0..30 {
                 let sk = SecretKey::random(&params, &mut rng);
-                let v = params.plaintext.random_vec(params.degree(), &mut rng);
+                let v = fhe_math::zq::Modulus::new(params.plaintext()).unwrap().random_vec(params.degree(), &mut rng);
                 let row_size = params.degree() >> 1;
 
                 let pt = Plaintext::try_encode(&v, Encoding::simd(), &params)?;

crates/fhe/src/bfv/keys/public_key.rs

@@ -188,7 +188,7 @@ mod tests {
                     let pk = PublicKey::new(&sk, &mut rng);
 
                     let pt = Plaintext::try_encode(
-                        &params.plaintext.random_vec(params.degree(), &mut rng),
+                        &fhe_math::zq::Modulus::new(params.plaintext()).unwrap().random_vec(params.degree(), &mut rng),
                         Encoding::poly_at_level(level),
                         &params,
                     )?;

crates/fhe/src/bfv/keys/secret_key.rs

@@ -1,6 +1,6 @@
 //! Secret keys for the BFV encryption scheme
 
-use crate::bfv::{BfvParameters, Ciphertext, Plaintext};
+use crate::bfv::{BfvParameters, Ciphertext, Plaintext, plaintext::PlaintextValues, parameters::PlaintextModulus};
 use crate::proto::bfv::SecretKey as SecretKeyProto;
 use crate::{Error, Result, SerializationError};
 use fhe_math::{
@@ -237,25 +237,56 @@ impl FheDecrypter<Plaintext, Ciphertext> for SecretKey {
             let ctx_lvl = self.par.context_level_at(ct.level).unwrap();
             let d = Zeroizing::new(c.scale(&ctx_lvl.cipher_plain_context.scaler)?);
 
-            // TODO: Can we handle plaintext moduli that are BigUint?
-            let v = Zeroizing::new(
-                Vec::<u64>::try_from(d.as_ref())?
-                    .iter_mut()
-                    .map(|vi| *vi + *self.par.plaintext)
-                    .collect_vec(),
-            );
-            let mut w = v[..self.par.degree()].to_vec();
-            let q = Modulus::new(self.par.moduli[0]).map_err(Error::MathError)?;
-            q.reduce_vec(&mut w);
-            self.par.plaintext.reduce_vec(&mut w);
-
-            let mut poly =
-                Poly::try_convert_from(&w, ct[0].ctx(), false, Representation::PowerBasis)?;
+            let value = match self.par.plaintext {
+                PlaintextModulus::Small(_) => {
+                    let mut v = Vec::<u64>::try_from(d.as_ref())?;
+                    let plaintext_modulus = self.par.plaintext();
+                    v.iter_mut().for_each(|vi| *vi += plaintext_modulus);
+                    let mut w = v[..self.par.degree()].to_vec();
+
+                    let q = Modulus::new(self.par.moduli[0]).map_err(Error::MathError)?;
+                    q.reduce_vec(&mut w);
+                    if let PlaintextModulus::Small(ref m) = self.par.plaintext {
+                        m.reduce_vec(&mut w);
+                    }
+                    PlaintextValues::Small(w.into_boxed_slice())
+                },
+                PlaintextModulus::Large(_) => {
+                    let v: Vec<BigUint> = Vec::<BigUint>::from(d.as_ref())
+                            .into_iter()
+                            .map(|vi| vi + self.par.plaintext_big())
+                            .collect_vec();
+
+                    let mut w = v[..self.par.degree()].to_vec();
+                    let q_poly = d.as_ref().ctx().modulus();
+                    w.iter_mut().for_each(|wi| *wi %= q_poly);
+
+                    self.par.plaintext.reduce_vec(&mut w);
+                    PlaintextValues::Large(w.into_boxed_slice())
+                }
+            };
+
+            let _poly_slice: &[BigUint] = match &value {
+                PlaintextValues::Small(_v) => {
+                    // This is inefficient but necessary if we want to call Poly::try_convert_from which expects &[BigUint] for Large
+                    // But Poly::try_convert_from can take &[u64].
+                    // Wait, we need to generate poly_ntt.
+                    // We can match again.
+                    &[] // dummy
+                },
+                PlaintextValues::Large(v) => v
+            };
+
+            let mut poly = match &value {
+                PlaintextValues::Small(v) => Poly::try_convert_from(v.as_ref(), ct[0].ctx(), false, Representation::PowerBasis)?,
+                PlaintextValues::Large(v) => Poly::try_convert_from(v.as_ref().as_ref(), ct[0].ctx(), false, Representation::PowerBasis)?
+            };
+
             poly.change_representation(Representation::Ntt);
 
             let pt = Plaintext {
                 par: self.par.clone(),
-                value: w.into_boxed_slice(),
+                value,
                 encoding: None,
                 poly_ntt: poly,
                 level: ct.level,
@@ -299,9 +330,10 @@ mod tests {
             for level in 0..params.max_level() {
                 for _ in 0..20 {
                     let sk = SecretKey::random(&params, &mut rng);
+                    let q = fhe_math::zq::Modulus::new(params.plaintext()).unwrap();
 
                     let pt = Plaintext::try_encode(
-                        &params.plaintext.random_vec(params.degree(), &mut rng),
+                        &q.random_vec(params.degree(), &mut rng),
                         Encoding::poly_at_level(level),
                         &params,
                     )?;

crates/fhe/src/bfv/mod.rs

@@ -25,6 +25,8 @@ pub(crate) use keys::KeySwitchingKey;
 pub use keys::{EvaluationKey, EvaluationKeyBuilder, PublicKey, RelinearizationKey, SecretKey};
 pub use ops::{Multiplicator, dot_product_scalar};
 pub use parameters::{BfvParameters, BfvParametersBuilder};
+pub(crate) use parameters::PlaintextModulus;
 pub use plaintext::Plaintext;
+pub(crate) use plaintext::PlaintextValues;
 pub use plaintext_vec::PlaintextVec;
 pub use rgsw_ciphertext::RGSWCiphertext;

crates/fhe/src/bfv/ops/dot_product.rs

@@ -180,14 +180,14 @@ mod tests {
             for size in 1..128 {
                 let ct = (0..size)
                     .map(|_| {
-                        let v = params.plaintext.random_vec(params.degree(), &mut rng);
+                        let v = fhe_math::zq::Modulus::new(params.plaintext()).unwrap().random_vec(params.degree(), &mut rng);
                         let pt = Plaintext::try_encode(&v, Encoding::simd(), &params).unwrap();
                         sk.try_encrypt(&pt, &mut rng).unwrap()
                     })
                     .collect_vec();
                 let pt = (0..size)
                     .map(|_| {
-                        let v = params.plaintext.random_vec(params.degree(), &mut rng);
+                        let v = fhe_math::zq::Modulus::new(params.plaintext()).unwrap().random_vec(params.degree(), &mut rng);
                         Plaintext::try_encode(&v, Encoding::simd(), &params).unwrap()
                     })
                     .collect_vec();