Extract out GLev encryption into its own operations (#382)

sunscreen_tfhe/src/entities/glev_ciphertext.rs

@@ -1,14 +1,16 @@
-use num::Complex;
+use num::{Complex, Zero};
 use serde::{Deserialize, Serialize};
 
-use crate::{dst::OverlaySize, GlweDef, GlweDimension, RadixCount, Torus, TorusOps};
+use crate::{
+    dst::OverlaySize, GlweDef, GlweDimension, RadixCount, RadixDecomposition, Torus, TorusOps,
+};
 
 use super::{
     GlevCiphertextFftRef, GlweCiphertextIterator, GlweCiphertextIteratorMut, GlweCiphertextRef,
 };
 
 dst! {
-    /// A GLEV ciphertext. For the FFT variant, see
+    /// A GLev ciphertext. For the FFT variant, see
     /// [`GlevCiphertextFft`](crate::entities::GlevCiphertextFft).
     GlevCiphertext,
     GlevCiphertextRef,
@@ -29,6 +31,20 @@ where
     }
 }
 
+impl<S> GlevCiphertext<S>
+where
+    S: TorusOps,
+{
+    /// Create a new zero GLev ciphertext with the given parameters.
+    pub fn new(params: &GlweDef, radix: &RadixDecomposition) -> Self {
+        let elems = GlevCiphertextRef::<S>::size((params.dim, radix.count));
+
+        Self {
+            data: avec![Torus::zero(); elems],
+        }
+    }
+}
+
 impl<S> GlevCiphertextRef<S>
 where
     S: TorusOps,
@@ -55,4 +71,12 @@ where
             i.fft(fft, params);
         }
     }
+
+    /// Assert that this entityt is valid.
+    pub fn assert_valid(&self, params: &GlweDef, radix: &RadixDecomposition) {
+        assert_eq!(
+            self.data.len(),
+            GlevCiphertextRef::<S>::size((params.dim, radix.count))
+        );
+    }
 }

sunscreen_tfhe/src/ops/encryption/ggsw_encryption.rs

@@ -1,19 +1,17 @@
-use num::Zero;
-
 use crate::{
     dst::FromMutSlice,
     entities::{GgswCiphertextRef, GlweCiphertextRef, GlweSecretKeyRef, Polynomial, PolynomialRef},
-    polynomial::{polynomial_external_mad, polynomial_scalar_mad, polynomial_scalar_mul},
+    ops::encryption::encrypt_glev_ciphertext_generic,
+    polynomial::polynomial_external_mad,
     scratch::allocate_scratch_ref,
     GlweDef, PlaintextBits, RadixDecomposition, Torus, TorusOps,
 };
 
 use super::{
-    decrypt_glwe_ciphertext, encrypt_glwe_ciphertext_secret,
-    trivially_encrypt_glwe_with_sk_argument,
+    decrypt_glwe_in_glev, encrypt_glwe_ciphertext_secret, trivially_encrypt_glwe_with_sk_argument,
 };
 
-/// Perform a ggsw encryption. This is generic in case a trivial GGSW encryption
+/// Perform a GGSW encryption. This is generic in case a trivial GGSW encryption
 /// is wanted (for example, for testing purposes).
 pub(crate) fn encrypt_ggsw_ciphertext_generic<S>(
     ggsw_ciphertext: &mut GgswCiphertextRef<S>,
@@ -34,7 +32,6 @@ pub(crate) fn encrypt_ggsw_ciphertext_generic<S>(
     let max_val = S::from_u64(0x1 << plaintext_bits.0);
     assert!(msg.coeffs().iter().all(|x| *x < max_val));
 
-    let decomposition_radix_log = radix.radix_log.0;
     let polynomial_degree = params.dim.polynomial_degree.0;
     let glwe_size = params.dim.size.0;
 
@@ -62,18 +59,7 @@ pub(crate) fn encrypt_ggsw_ciphertext_generic<S>(
             msg.as_torus()
         };
 
-        for (j, col) in row.glwe_ciphertexts_mut(params).enumerate() {
-            let mut scaled_msg = Polynomial::zero(polynomial_degree);
-
-            // The factor is q / B^{i+1}. Since B is a power of 2, this is equivalent to
-            // multiplying by 2^{log2(q) - log2(B) * (i + 1)}
-            let decomp_factor =
-                S::from_u64(0x1 << (S::BITS as usize - decomposition_radix_log * (j + 1)));
-
-            polynomial_scalar_mul(&mut scaled_msg, m_times_s, decomp_factor);
-
-            encrypt(col, &scaled_msg, glwe_secret_key, params);
-        }
+        encrypt_glev_ciphertext_generic(row, m_times_s, glwe_secret_key, params, radix, &encrypt);
     }
 }
 
@@ -130,63 +116,25 @@ pub fn encrypt_ggsw_ciphertext_scalar<S>(
     ggsw_ciphertext: &mut GgswCiphertextRef<S>,
     msg: S,
     glwe_secret_key: &GlweSecretKeyRef<S>,
-    glwe_def: &GlweDef,
+    params: &GlweDef,
     radix: &RadixDecomposition,
     plaintext_bits: PlaintextBits,
 ) where
     S: TorusOps,
 {
-    assert!(plaintext_bits.0 < S::BITS);
-    radix.assert_valid::<S>();
-    glwe_def.assert_valid();
-    glwe_secret_key.assert_valid(glwe_def);
-    ggsw_ciphertext.assert_valid(glwe_def, radix);
+    let polynomial_degree = params.dim.polynomial_degree.0;
 
-    let max_val = S::from_u64(0x1 << plaintext_bits.0);
-    assert!(msg < max_val);
+    let mut poly_msg = Polynomial::<S>::zero(polynomial_degree);
+    poly_msg.coeffs_mut()[0] = msg;
 
-    let decomposition_radix_log = radix.radix_log.0;
-    let polynomial_degree = glwe_def.dim.polynomial_degree.0;
-    let glwe_size = glwe_def.dim.size.0;
-
-    // k + 1 rows with l columns of glwe ciphertexts. Element (i,j) is a glwe encryption
-    // of -M/B^{i+1} * s_j, except for j=k+1, where it's simply an encryption of
-    // M/B^{j+1}
-    for (i, row) in ggsw_ciphertext.rows_mut(glwe_def, radix).enumerate() {
-        let mut m_times_s = Polynomial::<Torus<S>>::zero(polynomial_degree);
-        let m_times_s = if i < glwe_size {
-            let s = glwe_secret_key.s(glwe_def).nth(i).unwrap();
-            polynomial_scalar_mad(&mut m_times_s, s.as_torus(), msg);
-            &m_times_s
-        } else {
-            // Last row isn't multiplied by secret key.
-            m_times_s.clear();
-            m_times_s.coeffs_mut()[0] = Torus::from(msg);
-            &m_times_s
-        };
-
-        for (j, col) in row.glwe_ciphertexts_mut(glwe_def).enumerate() {
-            let mut scaled_msg = Polynomial::zero(polynomial_degree);
-            // The factor is q / B^{i+1}. Since B is a power of 2, this is equivalent to
-            // multiplying by 2^{log2(q) - log2(B) * (i + 1)}
-            let decomp_factor =
-                S::from_u64(0x1 << (S::BITS as usize - decomposition_radix_log * (j + 1)));
-
-            if i < glwe_size {
-                let decomp_factor = decomp_factor.wrapping_neg();
-
-                polynomial_scalar_mul(&mut scaled_msg, m_times_s, decomp_factor);
-            } else {
-                scaled_msg.coeffs_mut()[0] = Torus::from(msg.wrapping_mul(&decomp_factor));
-
-                for c in scaled_msg.coeffs_mut().iter_mut().skip(1) {
-                    *c = Torus::zero();
-                }
-            }
-
-            encrypt_glwe_ciphertext_secret(col, &scaled_msg, glwe_secret_key, glwe_def);
-        }
-    }
+    encrypt_ggsw_ciphertext(
+        ggsw_ciphertext,
+        &poly_msg,
+        glwe_secret_key,
+        params,
+        radix,
+        plaintext_bits,
+    )
 }
 
 fn decrypt_glwe_in_ggsw<S>(
@@ -201,27 +149,8 @@ fn decrypt_glwe_in_ggsw<S>(
 where
     S: TorusOps,
 {
-    let decomposition_radix_log = radix.radix_log.0;
-
-    // To decrypt a GGSW ciphertext, it suffices to decrypt the first GLWE ciphertext in
-    // the last row and divide by its decomposition factor.
     let glev = ggsw_ciphertext.rows(params, radix).nth(row)?;
-    let glwe = glev.glwe_ciphertexts(params).nth(column)?;
-
-    // Decrypt that specific GLWE ciphertext, which should have a message of
-    // q / beta ^ {column + 1} * SM, where SM is the message times the secret
-    // every row but the last (-SM) and M for the last row.
-    decrypt_glwe_ciphertext(msg, glwe, glwe_secret_key, params);
-
-    let mask = (0x1 << decomposition_radix_log) - 1;
-
-    for c in msg.coeffs_mut() {
-        let val = c.inner() >> (S::BITS as usize - decomposition_radix_log * (column + 1));
-        let r = (c.inner() >> (S::BITS as usize - decomposition_radix_log * (column + 1) - 1))
-            & S::from_u64(0x1);
-
-        *c = Torus::from((val + r) & S::from_u64(mask));
-    }
+    decrypt_glwe_in_glev(msg, glev, glwe_secret_key, params, radix, column)?;
 
     Some(())
 }
@@ -242,8 +171,11 @@ pub fn decrypt_ggsw_ciphertext<S>(
     ggsw_ciphertext.assert_valid(params, radix);
     glwe_secret_key.assert_valid(params);
 
+    // To decrypt a GGSW ciphertext, it suffices to decrypt the first GLWE
+    // ciphertext in the last row. We can decrypt any of the GLWE ciphertexts in
+    // the last row and divide them by their decomposition factor; we choose the
+    // first GLWE ciphertext.
     let row = params.dim.size.0;
-
     decrypt_glwe_in_ggsw(msg, ggsw_ciphertext, glwe_secret_key, params, radix, row, 0).unwrap();
 }
 

sunscreen_tfhe/src/ops/encryption/glev_encryption.rs

@@ -0,0 +1,271 @@
+use crate::{
+    dst::FromMutSlice,
+    entities::{GlevCiphertextRef, GlweCiphertextRef, GlweSecretKeyRef, Polynomial, PolynomialRef},
+    polynomial::polynomial_scalar_mul,
+    scratch::allocate_scratch_ref,
+    GlweDef, RadixDecomposition, Torus, TorusOps,
+};
+
+use super::{
+    decrypt_glwe_ciphertext, encrypt_glwe_ciphertext_secret,
+    trivially_encrypt_glwe_with_sk_argument,
+};
+
+/// Perform a GLev encryption. This is generic in case a trivial GLev encryption
+/// is wanted (for example, for testing purposes).
+pub(crate) fn encrypt_glev_ciphertext_generic<S>(
+    glev_ciphertext: &mut GlevCiphertextRef<S>,
+    msg: &PolynomialRef<Torus<S>>,
+    glwe_secret_key: &GlweSecretKeyRef<S>,
+    params: &GlweDef,
+    radix: &RadixDecomposition,
+    encrypt: impl Fn(
+        &mut GlweCiphertextRef<S>,
+        &PolynomialRef<Torus<S>>,
+        &GlweSecretKeyRef<S>,
+        &GlweDef,
+    ),
+) where
+    S: TorusOps,
+{
+    let decomposition_radix_log = radix.radix_log.0;
+    let polynomial_degree = params.dim.polynomial_degree.0;
+
+    for (j, glwe) in glev_ciphertext.glwe_ciphertexts_mut(params).enumerate() {
+        let mut scaled_msg = Polynomial::zero(polynomial_degree);
+
+        // The factor is q / B^{i+1}. Since B is a power of 2, this is equivalent to
+        // multiplying by 2^{log2(q) - log2(B) * (i + 1)}
+        let decomp_factor =
+            S::from_u64(0x1 << (S::BITS as usize - decomposition_radix_log * (j + 1)));
+
+        polynomial_scalar_mul(&mut scaled_msg, msg, decomp_factor);
+
+        encrypt(glwe, &scaled_msg, glwe_secret_key, params);
+    }
+}
+
+#[allow(dead_code)]
+/// Encrypt a GLev ciphertext with a given message polynomial and secret key.
+pub fn encrypt_glev_ciphertext<S>(
+    glev_ciphertext: &mut GlevCiphertextRef<S>,
+    msg: &PolynomialRef<Torus<S>>,
+    glwe_secret_key: &GlweSecretKeyRef<S>,
+    params: &GlweDef,
+    radix: &RadixDecomposition,
+) where
+    S: TorusOps,
+{
+    encrypt_glev_ciphertext_generic(
+        glev_ciphertext,
+        msg,
+        glwe_secret_key,
+        params,
+        radix,
+        encrypt_glwe_ciphertext_secret,
+    );
+}
+
+#[allow(dead_code)]
+/// Encrypt a GLev ciphertext with a given message polynomial and secret key.
+/// This is a trivial encryption that doesn't use the secret key and is not
+/// secure.
+pub fn trivially_encrypt_glev_ciphertext<S>(
+    glev_ciphertext: &mut GlevCiphertextRef<S>,
+    msg: &PolynomialRef<Torus<S>>,
+    params: &GlweDef,
+    radix: &RadixDecomposition,
+) where
+    S: TorusOps,
+{
+    allocate_scratch_ref!(trivial_key, GlweSecretKeyRef<S>, (params.dim));
+    trivial_key.clear();
+
+    encrypt_glev_ciphertext_generic(
+        glev_ciphertext,
+        msg,
+        trivial_key,
+        params,
+        radix,
+        trivially_encrypt_glwe_with_sk_argument,
+    );
+}
+
+pub(crate) fn decrypt_glwe_in_glev<S>(
+    msg: &mut PolynomialRef<Torus<S>>,
+    glev_ciphertext: &GlevCiphertextRef<S>,
+    glwe_secret_key: &GlweSecretKeyRef<S>,
+    params: &GlweDef,
+    radix: &RadixDecomposition,
+    index: usize,
+) -> Option<()>
+where
+    S: TorusOps,
+{
+    let decomposition_radix_log = radix.radix_log.0;
+
+    // To decrypt a GLev ciphertext, it suffices to decrypt the first GLWE ciphertext in
+    // and divide by its decomposition factor.
+    let glwe = glev_ciphertext.glwe_ciphertexts(params).nth(index)?;
+
+    // Decrypt that specific GLWE ciphertext, which should have a message of
+    // q / beta ^ {column + 1} * SM, where SM is the message times the secret
+    // every row but the last (-SM) and M for the last row.
+    decrypt_glwe_ciphertext(msg, glwe, glwe_secret_key, params);
+
+    let mask = (0x1 << decomposition_radix_log) - 1;
+
+    for c in msg.coeffs_mut() {
+        let val = c.inner() >> (S::BITS as usize - decomposition_radix_log * (index + 1));
+        let r = (c.inner() >> (S::BITS as usize - decomposition_radix_log * (index + 1) - 1))
+            & S::from_u64(0x1);
+
+        *c = Torus::from((val + r) & S::from_u64(mask));
+    }
+
+    Some(())
+}
+
+#[allow(dead_code)]
+/// Decrypt a GLev ciphertext with a given secret key.
+pub fn decrypt_glev_ciphertext<S>(
+    msg: &mut PolynomialRef<Torus<S>>,
+    glev_ciphertext: &GlevCiphertextRef<S>,
+    glwe_secret_key: &GlweSecretKeyRef<S>,
+    params: &GlweDef,
+    radix: &RadixDecomposition,
+) where
+    S: TorusOps,
+{
+    assert_eq!(msg.len(), params.dim.polynomial_degree.0);
+    params.assert_valid();
+    radix.assert_valid::<S>();
+    glev_ciphertext.assert_valid(params, radix);
+    glwe_secret_key.assert_valid(params);
+
+    decrypt_glwe_in_glev(msg, glev_ciphertext, glwe_secret_key, params, radix, 0).unwrap();
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{entities::GlevCiphertext, high_level::*};
+
+    use super::*;
+
+    fn encrypt_decrypt_glev_const_coeff() {
+        let params = &TEST_GLWE_DEF_1;
+        let radix = &TEST_RADIX;
+
+        let sk = keygen::generate_binary_glwe_sk(params);
+
+        let msg = 1u64;
+        let mut poly_msg = Polynomial::zero(params.dim.polynomial_degree.0);
+        poly_msg.coeffs_mut()[0] = msg;
+        let poly_msg = poly_msg.as_torus();
+
+        let mut glev_ciphertext = GlevCiphertext::new(params, radix);
+        let mut output_msg = Polynomial::zero(params.dim.polynomial_degree.0);
+
+        encrypt_glev_ciphertext(&mut glev_ciphertext, poly_msg, &sk, params, radix);
+        decrypt_glev_ciphertext(&mut output_msg, &glev_ciphertext, &sk, params, radix);
+
+        assert_eq!(output_msg.coeffs()[0], msg.into());
+
+        for c in output_msg.coeffs().iter().skip(1) {
+            assert_eq!(*c, 0.into());
+        }
+    }
+
+    #[test]
+    fn can_encrypt_decrypt_glev_const_coeff() {
+        for _ in 0..10 {
+            encrypt_decrypt_glev_const_coeff();
+        }
+    }
+
+    /// Test that each of the rows in the GLev ciphertext is a GLWE ciphertext that encodes the
+    /// appropriate message (usually the decomposed message times the secret key)
+    #[test]
+    fn can_decrypt_all_elements_glev() {
+        let params = TEST_GLWE_DEF_1;
+        let radix = TEST_RADIX;
+
+        let sk = keygen::generate_binary_glwe_sk(&params);
+
+        let coeffs = (0..params.dim.polynomial_degree.0 as u64)
+            .map(|x| x % 2)
+            .collect::<Vec<_>>();
+        let msg = Polynomial::new(&coeffs).as_torus().to_owned();
+
+        let mut ct = GlevCiphertext::new(&params, &radix);
+        encrypt_glev_ciphertext(&mut ct, &msg, &sk, &params, &radix);
+
+        let mut pt = Polynomial::zero(params.dim.polynomial_degree.0);
+        decrypt_glev_ciphertext(&mut pt, &ct, &sk, &params, &radix);
+
+        // Ensure that the basic decryption works.
+        assert_eq!(pt, msg.to_owned());
+
+        let n_glwes = ct.glwe_ciphertexts(&params).len();
+
+        // Beta
+        let decomposition_radix_log = radix.radix_log.0;
+
+        for i in 0..n_glwes {
+            let mut pt = Polynomial::zero(params.dim.polynomial_degree.0);
+            let mut scaled_msg = msg.to_owned();
+
+            let mask = (0x1 << decomposition_radix_log) - 1;
+
+            for c in scaled_msg.coeffs_mut() {
+                *c = Torus::from(c.inner() & mask);
+            }
+
+            decrypt_glwe_in_glev(&mut pt, &ct, &sk, &params, &radix, i).unwrap();
+
+            assert_eq!(pt, scaled_msg.to_owned());
+        }
+    }
+
+    #[test]
+    fn can_trivially_decrypy_glev() {
+        let params = TEST_GLWE_DEF_1;
+        let radix = TEST_RADIX;
+
+        let sk = keygen::generate_binary_glwe_sk(&params);
+
+        let coeffs = (0..params.dim.polynomial_degree.0 as u64)
+            .map(|x| x % 2)
+            .collect::<Vec<_>>();
+        let msg = Polynomial::new(&coeffs).as_torus().to_owned();
+
+        let mut ct = GlevCiphertext::new(&params, &radix);
+        trivially_encrypt_glev_ciphertext(&mut ct, &msg, &params, &radix);
+
+        let mut pt = Polynomial::zero(params.dim.polynomial_degree.0);
+        decrypt_glev_ciphertext(&mut pt, &ct, &sk, &params, &radix);
+
+        // Ensure that the basic decryption works.
+        assert_eq!(pt, msg);
+
+        let n_glwe = ct.glwe_ciphertexts(&params).len();
+
+        // Beta
+        let decomposition_radix_log = radix.radix_log.0;
+
+        for i in 0..n_glwe {
+            let mut pt = Polynomial::zero(params.dim.polynomial_degree.0);
+            let mut scaled_msg = msg.to_owned();
+
+            let mask = (0x1 << decomposition_radix_log) - 1;
+
+            for c in scaled_msg.coeffs_mut() {
+                *c = Torus::from(c.inner() & mask);
+            }
+
+            decrypt_glwe_in_glev(&mut pt, &ct, &sk, &params, &radix, i).unwrap();
+
+            assert_eq!(pt, scaled_msg);
+        }
+    }
+}

sunscreen_tfhe/src/ops/encryption/mod.rs

@@ -1,6 +1,9 @@
 mod lwe_encryption;
 pub use lwe_encryption::*;
 
+mod glev_encryption;
+pub use glev_encryption::*;
+
 mod glwe_encryption;
 pub use glwe_encryption::*;