chore(shortint): add constructor for compressed key switching key

zama-ai · Jul 19, 2024 · 5b37a83 · 5b37a83
1 parent c1fcd95
commit 5b37a83
Show file tree

Hide file tree

Showing 2 changed files with 111 additions and 0 deletions.
diff --git a/tfhe/src/shortint/key_switching_key/mod.rs b/tfhe/src/shortint/key_switching_key/mod.rs
@@ -733,4 +733,108 @@ impl CompressedKeySwitchingKey {
                 .map(CompressedServerKey::decompress),
         }
     }
+
+    /// Deconstruct a [`CompressedKeySwitchingKey`] into its constituents.
+    pub fn into_raw_parts(
+        self,
+    ) -> (
+        CompressedKeySwitchingKeyMaterial,
+        CompressedServerKey,
+        Option<CompressedServerKey>,
+    ) {
+        let Self {
+            key_switching_key_material,
+            dest_server_key,
+            src_server_key,
+        } = self;
+
+        (key_switching_key_material, dest_server_key, src_server_key)
+    }
+
+    /// Construct a [`CompressedKeySwitchingKey`] from its constituents.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the provided raw parts are not compatible with each other, i.e.:
+    ///
+    /// if the provided source [`CompressedServerKey`] ciphertext
+    /// [`LweDimension`](`crate::core_crypto::commons::parameters::LweDimension`) does not match the
+    /// input [`LweDimension`](`crate::core_crypto::commons::parameters::LweDimension`) of the
+    /// [`SeededLweKeyswitchKeyOwned`] in the provided [`CompressedKeySwitchingKeyMaterial`] or if
+    /// the provided destination [`CompressedServerKey`] ciphertext
+    /// [`LweDimension`](`crate::core_crypto::commons::parameters::LweDimension`) does not match
+    /// the output [`LweDimension`](`crate::core_crypto::commons::parameters::LweDimension`) of
+    /// the [`SeededLweKeyswitchKeyOwned`] in the provided [`CompressedKeySwitchingKeyMaterial`].
+    pub fn from_raw_parts(
+        key_switching_key_material: CompressedKeySwitchingKeyMaterial,
+        dest_server_key: CompressedServerKey,
+        src_server_key: Option<CompressedServerKey>,
+    ) -> Self {
+        match src_server_key {
+            Some(ref src_server_key) => {
+                let src_lwe_dimension = src_server_key.ciphertext_lwe_dimension();
+
+                assert_eq!(
+                    src_lwe_dimension,
+                    key_switching_key_material
+                        .key_switching_key
+                        .input_key_lwe_dimension(),
+                    "Mismatch between the source CompressedServerKey ciphertext LweDimension ({:?}) \
+                    and the SeededLweKeyswitchKey input LweDimension ({:?})",
+                    src_lwe_dimension,
+                    key_switching_key_material
+                        .key_switching_key
+                        .input_key_lwe_dimension(),
+                );
+
+                assert_eq!(
+                    src_server_key.ciphertext_modulus, dest_server_key.ciphertext_modulus,
+                    "Mismatch between the source CompressedServerKey CiphertextModulus ({:?}) \
+                    and the destination CompressedServerKey CiphertextModulus ({:?})",
+                    src_server_key.ciphertext_modulus, dest_server_key.ciphertext_modulus,
+                );
+            }
+            None => assert!(
+                key_switching_key_material.cast_rshift >= 0,
+                "Trying to build a shortint::CompressedKeySwitchingKey with a negative cast_rshift \
+                without providing a source CompressedServerKey, this is not supported"
+            ),
+        }
+
+        let dst_lwe_dimension = match key_switching_key_material.destination_key {
+            EncryptionKeyChoice::Big => dest_server_key.bootstrapping_key.output_lwe_dimension(),
+            EncryptionKeyChoice::Small => dest_server_key.bootstrapping_key.input_lwe_dimension(),
+        };
+
+        assert_eq!(
+            dst_lwe_dimension,
+            key_switching_key_material
+                .key_switching_key
+                .output_key_lwe_dimension(),
+            "Mismatch between the destination CompressedServerKey ciphertext LweDimension ({:?}) \
+            and the SeededLweKeyswitchKey output LweDimension ({:?})",
+            dst_lwe_dimension,
+            key_switching_key_material
+                .key_switching_key
+                .output_key_lwe_dimension(),
+        );
+        assert_eq!(
+            key_switching_key_material
+                .key_switching_key
+                .ciphertext_modulus(),
+            dest_server_key.ciphertext_modulus,
+            "Mismatch between the SeededLweKeyswitchKey CiphertextModulus ({:?}) \
+            and the destination CompressedServerKey CiphertextModulus ({:?})",
+            key_switching_key_material
+                .key_switching_key
+                .ciphertext_modulus(),
+            dest_server_key.ciphertext_modulus,
+        );
+
+        Self {
+            key_switching_key_material,
+            dest_server_key,
+            src_server_key,
+        }
+    }
 }
diff --git a/tfhe/src/shortint/server_key/compressed.rs b/tfhe/src/shortint/server_key/compressed.rs
@@ -395,4 +395,11 @@ impl CompressedServerKey {
             engine.new_compressed_server_key_with_max_degree(cks, max_degree)
         })
     }
+
+    pub fn ciphertext_lwe_dimension(&self) -> LweDimension {
+        match self.pbs_order {
+            PBSOrder::KeyswitchBootstrap => self.key_switching_key.input_key_lwe_dimension(),
+            PBSOrder::BootstrapKeyswitch => self.key_switching_key.output_key_lwe_dimension(),
+        }
+    }
 }