dont_be_mean.patch

diff --git a/tfhe/src/core_crypto/algorithms/lwe_keyswitch.rs b/tfhe/src/core_crypto/algorithms/lwe_keyswitch.rs
index 4346f11..b8d756f 100644
--- a/tfhe/src/core_crypto/algorithms/lwe_keyswitch.rs
+++ b/tfhe/src/core_crypto/algorithms/lwe_keyswitch.rs
@@ -4,7 +4,7 @@
 use crate::core_crypto::algorithms::slice_algorithms::*;
 use crate::core_crypto::commons::ciphertext_modulus::CiphertextModulusKind;
 use crate::core_crypto::commons::math::decomposition::{
-    SignedDecomposer, SignedDecomposerNonNative,
+    SignedDecomposer, SignedDecomposerNonNative, SignedDecompositionIter
 };
 use crate::core_crypto::commons::parameters::{
     DecompositionBaseLog, DecompositionLevelCount, ThreadCount,
@@ -186,8 +186,43 @@ pub fn keyswitch_lwe_ciphertext_native_mod_compatible<Scalar, KSKCont, InputCont
     // Clear the output ciphertext, as it will get updated gradually
     output_lwe_ciphertext.as_mut().fill(Scalar::ZERO);
 
+    // Don't be mean: Changes
+    let mut correction_term = Scalar::ZERO;
+    let precision = lwe_keyswitch_key.decomposition_base_log().0 * lwe_keyswitch_key.decomposition_level_count().0;
+    let rounding_mask = Scalar::ONE << (Scalar::BITS - precision - 1);
+    let selection_mask = Scalar::wrapping_neg(Scalar::ONE << (Scalar::BITS - precision));
+    for (keyswitch_key_block, &input_mask_element) in lwe_keyswitch_key
+        .iter()
+        .zip(input_lwe_ciphertext.get_mask().as_ref())
+    {
+        let rounded_to_floor = (input_mask_element).wrapping_add(rounding_mask);
+        let rounded = rounded_to_floor.bitand(selection_mask);
+        let rounding_error = (input_mask_element).wrapping_sub(rounded);
+        correction_term = correction_term.wrapping_add(rounding_error);
+
+        let decomposition_iter = SignedDecompositionIter::new(
+            rounded >> (Scalar::BITS - precision),
+            lwe_keyswitch_key.decomposition_base_log(),
+            lwe_keyswitch_key.decomposition_level_count(),
+        );
+        // Loop over the levels
+        for (level_key_ciphertext, decomposed) in keyswitch_key_block.iter().zip(decomposition_iter)
+        {
+            slice_wrapping_sub_scalar_mul_assign(
+                output_lwe_ciphertext.as_mut(),
+                level_key_ciphertext.as_ref(),
+                decomposed.value(),
+            );
+        }
+    }
+
+    let msb = Scalar::ONE << (Scalar::BITS - 1);
+    correction_term = (correction_term >> 1).wrapping_add(correction_term.bitand(msb));
+
     // Copy the input body to the output ciphertext
-    *output_lwe_ciphertext.get_mut_body().data = *input_lwe_ciphertext.get_body().data;
+    *output_lwe_ciphertext.get_mut_body().data = ((*output_lwe_ciphertext.get_body().data)
+                                                    .wrapping_add(*input_lwe_ciphertext.get_body().data))
+                                                    .wrapping_sub(correction_term);
 
     // If the moduli are not the same, we need to round the body in the output ciphertext
     if output_ciphertext_modulus != input_ciphertext_modulus
@@ -203,27 +238,6 @@ pub fn keyswitch_lwe_ciphertext_native_mod_compatible<Scalar, KSKCont, InputCont
             output_decomposer.closest_representable(*output_lwe_ciphertext.get_mut_body().data);
     }
 
-    // We instantiate a decomposer
-    let decomposer = SignedDecomposer::new(
-        lwe_keyswitch_key.decomposition_base_log(),
-        lwe_keyswitch_key.decomposition_level_count(),
-    );
-
-    for (keyswitch_key_block, &input_mask_element) in lwe_keyswitch_key
-        .iter()
-        .zip(input_lwe_ciphertext.get_mask().as_ref())
-    {
-        let decomposition_iter = decomposer.decompose(input_mask_element);
-        // Loop over the levels
-        for (level_key_ciphertext, decomposed) in keyswitch_key_block.iter().zip(decomposition_iter)
-        {
-            slice_wrapping_sub_scalar_mul_assign(
-                output_lwe_ciphertext.as_mut(),
-                level_key_ciphertext.as_ref(),
-                decomposed.value(),
-            );
-        }
-    }
 }
 
 /// Specialized implementation of an LWE keyswitch when inputs have non power of two moduli.
diff --git a/tfhe/src/core_crypto/fft_impl/common.rs b/tfhe/src/core_crypto/fft_impl/common.rs
index 70f4d96..f85d90e 100644
--- a/tfhe/src/core_crypto/fft_impl/common.rs
+++ b/tfhe/src/core_crypto/fft_impl/common.rs
@@ -14,6 +14,14 @@ pub fn pbs_modulus_switch<Scalar: UnsignedInteger + CastInto<usize>>(
     modulus_switch(input, polynomial_size.to_blind_rotation_input_modulus_log()).cast_into()
 }
 
+// Don't be mean: Changes
+pub fn pbs_modulus_switch_comp<Scalar: UnsignedInteger + CastInto<usize>>(
+    input: Scalar,
+    polynomial_size: PolynomialSize,
+) -> Scalar {
+    modulus_switch(input, polynomial_size.to_blind_rotation_input_modulus_log())
+}
+
 pub fn modulus_switch<Scalar: UnsignedInteger>(
     input: Scalar,
     log_modulus: CiphertextModulusLog,
diff --git a/tfhe/src/core_crypto/fft_impl/fft64/crypto/bootstrap.rs b/tfhe/src/core_crypto/fft_impl/fft64/crypto/bootstrap.rs
index 497efab..234a6ba 100644
--- a/tfhe/src/core_crypto/fft_impl/fft64/crypto/bootstrap.rs
+++ b/tfhe/src/core_crypto/fft_impl/fft64/crypto/bootstrap.rs
@@ -16,7 +16,7 @@ use crate::core_crypto::commons::traits::{
 };
 use crate::core_crypto::commons::utils::izip;
 use crate::core_crypto::entities::*;
-use crate::core_crypto::fft_impl::common::{pbs_modulus_switch, FourierBootstrapKey};
+use crate::core_crypto::fft_impl::common::{pbs_modulus_switch, pbs_modulus_switch_comp, FourierBootstrapKey};
 use crate::core_crypto::fft_impl::fft64::math::fft::par_convert_polynomials_list_to_fourier;
 use crate::core_crypto::prelude::{CiphertextCount, CiphertextModulus, ContainerMut};
 use aligned_vec::{avec, ABox, CACHELINE_ALIGN};
@@ -285,7 +285,7 @@ impl FourierLweBootstrapKeyView<'_> {
     // CastInto required for PBS modulus switch which returns a usize
     pub fn blind_rotate_assign<InputScalar, OutputScalar>(
         self,
-        mut lut: GlweCiphertextMutView<'_, OutputScalar>,
+        lut: GlweCiphertextMutView<'_, OutputScalar>,
         lwe: LweCiphertextView<'_, InputScalar>,
         fft: FftView<'_>,
         stack: &mut PodStack,
@@ -298,30 +298,26 @@ impl FourierLweBootstrapKeyView<'_> {
         let lut_poly_size = lut.polynomial_size();
         let ciphertext_modulus = lut.ciphertext_modulus();
         assert!(ciphertext_modulus.is_compatible_with_native_modulus());
-        let monomial_degree = MonomialDegree(pbs_modulus_switch(*lwe_body.data, lut_poly_size));
-
-        lut.as_mut_polynomial_list()
-            .iter_mut()
-            .for_each(|mut poly| {
-                let (tmp_poly, _) = stack.make_aligned_raw(poly.as_ref().len(), CACHELINE_ALIGN);
-
-                let mut tmp_poly = Polynomial::from_container(&mut *tmp_poly);
-                tmp_poly.as_mut().copy_from_slice(poly.as_ref());
-                polynomial_wrapping_monic_monomial_div(&mut poly, &tmp_poly, monomial_degree);
-            });
 
+        // Don't be mean: changes
+        //first handle a_i
         // We initialize the ct_0 used for the successive cmuxes
         let mut ct0 = lut;
         let (ct1, stack) = stack.make_aligned_raw(ct0.as_ref().len(), CACHELINE_ALIGN);
         let mut ct1 =
             GlweCiphertextMutView::from_container(&mut *ct1, lut_poly_size, ciphertext_modulus);
+        let mut correction_term = InputScalar::ZERO;
+        let shift = InputScalar::BITS - lut_poly_size.to_blind_rotation_input_modulus_log().0;
+
 
         for (lwe_mask_element, bootstrap_key_ggsw) in
             izip!(lwe_mask.as_ref().iter(), self.into_ggsw_iter())
         {
             if *lwe_mask_element != InputScalar::ZERO {
+                let monomial_degree_to_cast = pbs_modulus_switch_comp(*lwe_mask_element, lut_poly_size);
+                correction_term = correction_term.wrapping_add((*lwe_mask_element).wrapping_sub(monomial_degree_to_cast << shift));
                 let monomial_degree =
-                    MonomialDegree(pbs_modulus_switch(*lwe_mask_element, lut_poly_size));
+                    MonomialDegree(monomial_degree_to_cast.cast_into());
 
                 // we effectively inline the body of cmux here, merging the initial subtraction
                 // operation with the monic polynomial multiplication, then performing the external
@@ -352,6 +348,25 @@ impl FourierLweBootstrapKeyView<'_> {
             }
         }
 
+        //handle b in the end
+        let msb = InputScalar::ONE << (InputScalar::BITS - 1);
+        let corrected_body = (*lwe_body.data).wrapping_sub((correction_term >> 1).wrapping_add(correction_term.bitand(msb)));
+
+        let body_to_cast = pbs_modulus_switch_comp(corrected_body, lut_poly_size);
+
+        let monomial_degree = MonomialDegree(body_to_cast.cast_into());
+
+        ct0.as_mut_polynomial_list()
+            .iter_mut()
+            .for_each(|mut poly| {
+                let (tmp_poly, _) = stack
+                    .make_aligned_raw(poly.as_ref().len(), CACHELINE_ALIGN);
+
+                let mut tmp_poly = Polynomial::from_container(&mut *tmp_poly);
+                tmp_poly.as_mut().copy_from_slice(poly.as_ref());
+                polynomial_wrapping_monic_monomial_div(&mut poly, &tmp_poly, monomial_degree);
+            });
+        
         if !ciphertext_modulus.is_native_modulus() {
             // When we convert back from the fourier domain, integer values will contain up to 53
             // MSBs with information. In our representation of power of 2 moduli < native modulus we