fix(shortint): decrypt with the intermediate key for the standard ap

tfhe/src/shortint/client_key/atomic_pattern/ks32.rs

@@ -9,7 +9,7 @@ use crate::shortint::backward_compatibility::client_key::atomic_pattern::KS32Ato
 use crate::shortint::client_key::{GlweSecretKeyOwned, LweSecretKeyOwned, LweSecretKeyView};
 use crate::shortint::engine::ShortintEngine;
 use crate::shortint::parameters::{DynamicDistribution, KeySwitch32PBSParameters};
-use crate::shortint::{AtomicPatternKind, ShortintParameterSet};
+use crate::shortint::{AtomicPatternKind, EncryptionKeyChoice, ShortintParameterSet};
 
 use super::EncryptionAtomicPattern;
 
@@ -158,6 +158,14 @@ impl EncryptionAtomicPattern for KS32AtomicPatternClientKey {
         self.glwe_secret_key.as_lwe_secret_key()
     }
 
+    fn intermediate_encryption_key(&self) -> LweSecretKeyView<'_, u64> {
+        panic!("KS32 AP does not support decrypting with the intermediate encryption key")
+    }
+
+    fn encryption_key_choice(&self) -> EncryptionKeyChoice {
+        self.parameters.encryption_key_choice()
+    }
+
     fn encryption_noise(&self) -> DynamicDistribution<u64> {
         // The KS32 atomic pattern is only supported with the KsPbs order
         self.parameters.glwe_noise_distribution()

tfhe/src/shortint/client_key/atomic_pattern/mod.rs

@@ -7,7 +7,9 @@ use tfhe_versionable::Versionize;
 use crate::shortint::backward_compatibility::client_key::atomic_pattern::AtomicPatternClientKeyVersions;
 use crate::shortint::engine::ShortintEngine;
 use crate::shortint::parameters::DynamicDistribution;
-use crate::shortint::{AtomicPatternKind, AtomicPatternParameters, ShortintParameterSet};
+use crate::shortint::{
+    AtomicPatternKind, AtomicPatternParameters, EncryptionKeyChoice, ShortintParameterSet,
+};
 
 use super::{LweSecretKeyOwned, LweSecretKeyView};
 
@@ -26,6 +28,14 @@ pub trait EncryptionAtomicPattern {
     /// The secret key used for encryption
     fn encryption_key(&self) -> LweSecretKeyView<'_, u64>;
 
+    /// The kind of secret key used for encryption
+    fn encryption_key_choice(&self) -> EncryptionKeyChoice;
+
+    /// Encryption key used for ciphertexts "in the middle" of the atomic pattern
+    ///
+    /// For KS-PBS this is the small key, for PBS-KS this is the big key
+    fn intermediate_encryption_key(&self) -> LweSecretKeyView<'_, u64>;
+
     /// The noise distribution used for encryption
     fn encryption_noise(&self) -> DynamicDistribution<u64>;
 
@@ -49,6 +59,14 @@ impl<T: EncryptionAtomicPattern> EncryptionAtomicPattern for &T {
         (*self).encryption_key()
     }
 
+    fn intermediate_encryption_key(&self) -> LweSecretKeyView<'_, u64> {
+        (*self).intermediate_encryption_key()
+    }
+
+    fn encryption_key_choice(&self) -> EncryptionKeyChoice {
+        (*self).encryption_key_choice()
+    }
+
     fn encryption_noise(&self) -> DynamicDistribution<u64> {
         (*self).encryption_noise()
     }
@@ -119,6 +137,20 @@ impl EncryptionAtomicPattern for AtomicPatternClientKey {
         }
     }
 
+    fn intermediate_encryption_key(&self) -> LweSecretKeyView<'_, u64> {
+        match self {
+            Self::Standard(ap) => ap.intermediate_encryption_key(),
+            Self::KeySwitch32(ap) => ap.intermediate_encryption_key(),
+        }
+    }
+
+    fn encryption_key_choice(&self) -> EncryptionKeyChoice {
+        match self {
+            Self::Standard(ap) => ap.encryption_key_choice(),
+            Self::KeySwitch32(ap) => ap.encryption_key_choice(),
+        }
+    }
+
     fn encryption_noise(&self) -> DynamicDistribution<u64> {
         match self {
             Self::Standard(ap) => ap.encryption_noise(),

tfhe/src/shortint/client_key/atomic_pattern/standard.rs

@@ -254,14 +254,26 @@ impl EncryptionAtomicPattern for StandardAtomicPatternClientKey {
     }
 
     fn encryption_key(&self) -> LweSecretKeyView<'_, u64> {
-        match self.parameters.encryption_key_choice() {
+        match self.encryption_key_choice() {
             EncryptionKeyChoice::Big => self.large_lwe_secret_key(),
             EncryptionKeyChoice::Small => self.small_lwe_secret_key(),
         }
     }
 
+    fn intermediate_encryption_key(&self) -> LweSecretKeyView<'_, u64> {
+        // Use the opposite key as the one used for encryption
+        match self.encryption_key_choice() {
+            EncryptionKeyChoice::Big => self.small_lwe_secret_key(),
+            EncryptionKeyChoice::Small => self.large_lwe_secret_key(),
+        }
+    }
+
+    fn encryption_key_choice(&self) -> EncryptionKeyChoice {
+        self.parameters.encryption_key_choice()
+    }
+
     fn encryption_noise(&self) -> DynamicDistribution<u64> {
-        match self.parameters.encryption_key_choice() {
+        match self.encryption_key_choice() {
             EncryptionKeyChoice::Big => self.parameters.glwe_noise_distribution(),
             EncryptionKeyChoice::Small => self.parameters.lwe_noise_distribution(),
         }

tfhe/src/shortint/client_key/mod.rs

@@ -429,7 +429,13 @@ impl<AP: EncryptionAtomicPattern> GenericClientKey<AP> {
     /// assert!((noise as i64).abs() < delta as i64 / 2);
     /// ```
     pub fn decrypt_no_decode(&self, ct: &Ciphertext) -> Plaintext<u64> {
-        let lwe_decryption_key = self.encryption_key();
+        let lwe_decryption_key = if ct.atomic_pattern.pbs_order()
+            == self.atomic_pattern.encryption_key_choice().into_pbs_order()
+        {
+            self.encryption_key()
+        } else {
+            self.atomic_pattern.intermediate_encryption_key()
+        };
 
         decrypt_lwe_ciphertext(&lwe_decryption_key, &ct.ct)
     }