Implement scheme switching

sunscreen_tfhe/benches/ops.rs

@@ -6,12 +6,15 @@ use criterion::{
 
 use sunscreen_tfhe::{
     entities::{
-        GgswCiphertext, GgswCiphertextFft, GlweCiphertext, Polynomial, PolynomialRef,
-        PublicFunctionalKeyswitchKey, UnivariateLookupTable,
+        GgswCiphertext, GgswCiphertextFft, GlevCiphertext, GlweCiphertext, Polynomial,
+        PolynomialRef, PublicFunctionalKeyswitchKey, SchemeSwitchKey, SchemeSwitchKeyFft,
+        UnivariateLookupTable,
     },
     high_level::{self, *},
     ops::{
-        bootstrapping::circuit_bootstrap,
+        bootstrapping::{circuit_bootstrap, generate_scheme_switch_key},
+        encryption::encrypt_glev_ciphertext,
+        fft_ops::scheme_switch_fft,
         keyswitch::public_functional_keyswitch::{
             generate_public_functional_keyswitch_key, public_functional_keyswitch,
         },
@@ -234,6 +237,53 @@ fn circuit_bootstrapping(c: &mut Criterion) {
     });
 }
 
+fn scheme_switch(c: &mut Criterion) {
+    let ss_radix = RadixDecomposition {
+        count: RadixCount(2),
+        radix_log: RadixLog(19),
+    };
+    let ggsw_radix = RadixDecomposition {
+        count: RadixCount(1),
+        radix_log: RadixLog(11),
+    };
+
+    let params = GLWE_1_1024_80;
+    let polynomial_degree = params.dim.polynomial_degree.0;
+
+    // Generate the keys
+    let sk = keygen::generate_binary_glwe_sk(&params);
+
+    let mut ssk = SchemeSwitchKey::new(&params, &ss_radix);
+    generate_scheme_switch_key(&mut ssk, &sk, &params, &ss_radix);
+
+    let mut ssk_fft = SchemeSwitchKeyFft::new(&params, &ss_radix);
+    ssk.fft(&mut ssk_fft, &params, &ss_radix);
+
+    // Generate the message 1
+    let mut m_coeffs = vec![Torus::from(0u64); polynomial_degree];
+    m_coeffs[0] = Torus::from(1);
+    let m = Polynomial::new(&m_coeffs);
+
+    // Encrypt the message
+    let mut glev_ciphertext = GlevCiphertext::new(&params, &ggsw_radix);
+    encrypt_glev_ciphertext(&mut glev_ciphertext, &m, &sk, &params, &ggsw_radix);
+
+    let mut ggsw_fft = GgswCiphertextFft::new(&params, &ggsw_radix);
+
+    c.bench_function("Scheme Switch", |b| {
+        b.iter(|| {
+            scheme_switch_fft(
+                &mut ggsw_fft,
+                &glev_ciphertext,
+                &ssk_fft,
+                &params,
+                &ggsw_radix,
+                &ss_radix,
+            );
+        });
+    });
+}
+
 fn keygen(c: &mut Criterion) {
     c.bench_function("LWE Secret keygen", |b| {
         b.iter(|| {
@@ -347,6 +397,7 @@ criterion_group!(
     cmux,
     programmable_bootstrapping,
     circuit_bootstrapping,
+    scheme_switch,
     keygen,
     public_functional_keyswitching
 );

sunscreen_tfhe/src/entities/glwe_ciphertext.rs

@@ -157,4 +157,104 @@ where
             GlweCiphertextRef::<S>::size(params.dim)
         )
     }
+
+    /// Sets all coefficients of the polynomial at the specified index in the
+    /// GLWE ciphertext's mask (a) to zero.
+    ///
+    /// # Arguments
+    /// * `index` - The index of the polynomial in the mask to zero out
+    /// * `params` - The GLWE parameters defining the ciphertext structure
+    ///
+    /// # Panics
+    /// Panics if `index` is greater than or equal to the mask dimension (number of polynomials in a).
+    ///
+    /// # See also
+    /// * [`fill_a_at_index`] - Fills a polynomial at a specific index with provided coefficients
+    /// * [`zero_a_except_at_index`] - Zeros out all polynomials except at the specified index
+    #[allow(unused)]
+    pub(crate) fn zero_out_a_at_index(&mut self, index: usize, params: &GlweDef) {
+        assert!(index < params.dim.size.0, "index out of bounds");
+        self.a_mut(params)
+            .nth(index)
+            .unwrap()
+            .coeffs_mut()
+            .fill(Torus::from(<S as num::Zero>::zero()));
+    }
+
+    /// Copies the coefficients from the provided polynomial into the GLWE
+    /// ciphertext's mask (a) at the specified index.
+    ///
+    /// # Arguments
+    /// * `polynomial` - The source polynomial whose coefficients will be copied
+    /// * `index` - The index in the mask where the polynomial should be placed
+    /// * `params` - The GLWE parameters defining the ciphertext structure
+    ///
+    /// # Panics
+    /// * Panics if `index` is greater than or equal to the mask dimension
+    /// * Panics if the provided polynomial's length doesn't match the polynomial degree in params
+    ///
+    /// # See also
+    /// * [`zero_out_a_at_index`] - Zeros out a polynomial at a specific index
+    /// * [`zero_a_except_at_index`] - Zeros out all polynomials except at the specified index
+    pub(crate) fn fill_a_at_index(
+        &mut self,
+        polynomial: &PolynomialRef<Torus<S>>,
+        index: usize,
+        params: &GlweDef,
+    ) {
+        assert!(index < params.dim.size.0, "Index out of bounds");
+        assert_eq!(
+            polynomial.len(),
+            params.dim.polynomial_degree.0,
+            "Polynomial size mismatch"
+        );
+
+        self.a_mut(params)
+            .nth(index)
+            .unwrap()
+            .coeffs_mut()
+            .copy_from_slice(polynomial.coeffs());
+    }
+
+    /// Copies the coefficients from the provided polynomial into the GLWE
+    /// ciphertext's mask (a) at the specified index while setting all other
+    /// polynomials in the mask to zero.
+    ///
+    /// # Arguments
+    /// * `polynomial` - The source polynomial whose coefficients will be copied
+    /// * `index` - The index where the polynomial should be placed while zeroing out others
+    /// * `params` - The GLWE parameters defining the ciphertext structure
+    ///
+    /// # Panics
+    /// * Panics if `index` is greater than or equal to the mask dimension
+    /// * Panics if the provided polynomial's length doesn't match the polynomial degree in params
+    ///
+    /// # See also
+    /// * [`zero_out_a_at_index`] - Zeros out a polynomial at a specific index
+    /// * [`fill_a_at_index`] - Fills a polynomial at a specific index with provided coefficients
+    pub(crate) fn zero_a_except_at_index(
+        &mut self,
+        polynomial: &PolynomialRef<Torus<S>>,
+        index: usize,
+        params: &GlweDef,
+    ) {
+        assert!(index < params.dim.size.0, "Index out of bounds");
+        assert_eq!(
+            polynomial.len(),
+            params.dim.polynomial_degree.0,
+            "Polynomial size mismatch"
+        );
+
+        // Get iterator over all polynomials in the mask
+        for (i, poly) in self.a_mut(params).enumerate() {
+            if i == index {
+                // Fill with provided polynomial at specified index
+                poly.coeffs_mut().copy_from_slice(polynomial.coeffs());
+            } else {
+                // Zero out all other positions
+                poly.coeffs_mut()
+                    .fill(Torus::from(<S as num::Zero>::zero()));
+            }
+        }
+    }
 }

sunscreen_tfhe/src/entities/mod.rs

@@ -69,3 +69,9 @@ pub use polynomial_fft::*;
 
 mod polynomial_list;
 pub use polynomial_list::*;
+
+mod scheme_switch_key;
+pub use scheme_switch_key::*;
+
+mod scheme_switch_key_fft;
+pub use scheme_switch_key_fft::*;

sunscreen_tfhe/src/entities/polynomial.rs

@@ -365,6 +365,20 @@ where
     }
 }
 
+impl<T: PartialEq + Clone> PartialEq for PolynomialRef<T> {
+    fn eq(&self, other: &Self) -> bool {
+        self.coeffs() == other.coeffs()
+    }
+}
+
+impl<T: std::fmt::Debug + Clone> std::fmt::Debug for PolynomialRef<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("PolynomialRef")
+            .field("coeffs", &self.coeffs())
+            .finish()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::ops::Deref;