1use poulpy_hal::{
  2    api::{ScratchAvailable, SvpPPolBytesOf, VecZnxAutomorphism, VecZnxDftBytesOf, VecZnxNormalizeTmpBytes},
  3    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  4    source::Source,
  5};
  6
  7use crate::{
  8    encryption::compressed::gglwe_ksk::GGLWEKeyCompressedEncryptSk,
  9    layouts::{
 10        GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, LWEInfos,
 11        compressed::{AutomorphismKeyCompressed, AutomorphismKeyCompressedToMut, GLWESwitchingKeyCompressed},
 12    },
 13};
 14
 15impl AutomorphismKeyCompressed<Vec<u8>> {
 16    pub fn encrypt_sk_scratch_space<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
 17    where
 18        A: GGLWEInfos,
 19        Module<B>: VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + SvpPPolBytesOf,
 20    {
 21        assert_eq!(module.n() as u32, infos.n());
 22        GLWESwitchingKeyCompressed::encrypt_sk_scratch_space(module, infos) + GLWESecret::bytes_of(infos.n(), infos.rank_out())
 23    }
 24}
 25
 26pub trait GGLWEAutomorphismKeyCompressedEncryptSk<B: Backend> {
 27    fn gglwe_automorphism_key_compressed_encrypt_sk<R, S>(
 28        &self,
 29        res: &mut R,
 30        p: i64,
 31        sk: &S,
 32        seed_xa: [u8; 32],
 33        source_xe: &mut Source,
 34        scratch: &mut Scratch<B>,
 35    ) where
 36        R: AutomorphismKeyCompressedToMut,
 37        S: GLWESecretToRef;
 38}
 39
 40impl<B: Backend> GGLWEAutomorphismKeyCompressedEncryptSk<B> for Module<B>
 41where
 42    Module<B>: GGLWEKeyCompressedEncryptSk<B> + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + SvpPPolBytesOf + VecZnxAutomorphism,
 43    Scratch<B>: ScratchAvailable,
 44{
 45    fn gglwe_automorphism_key_compressed_encrypt_sk<R, S>(
 46        &self,
 47        res: &mut R,
 48        p: i64,
 49        sk: &S,
 50        seed_xa: [u8; 32],
 51        source_xe: &mut Source,
 52        scratch: &mut Scratch<B>,
 53    ) where
 54        R: AutomorphismKeyCompressedToMut,
 55        S: GLWESecretToRef,
 56    {
 57        let res: &mut AutomorphismKeyCompressed<&mut [u8]> = &mut res.to_mut();
 58        let sk: &GLWESecret<&[u8]> = &sk.to_ref();
 59
 60        #[cfg(debug_assertions)]
 61        {
 62            assert_eq!(res.n(), sk.n());
 63            assert_eq!(res.rank_out(), res.rank_in());
 64            assert_eq!(sk.rank(), res.rank_out());
 65            assert!(
 66                scratch.available() >= AutomorphismKeyCompressed::encrypt_sk_scratch_space(self, res),
 67                "scratch.available(): {} < AutomorphismKey::encrypt_sk_scratch_space: {}",
 68                scratch.available(),
 69                AutomorphismKeyCompressed::encrypt_sk_scratch_space(self, res)
 70            )
 71        }
 72
 73        let (mut sk_out, scratch_1) = scratch.take_glwe_secret(sk.n(), sk.rank());
 74
 75        {
 76            (0..res.rank_out().into()).for_each(|i| {
 77                self.vec_znx_automorphism(
 78                    self.galois_element_inv(p),
 79                    &mut sk_out.data.as_vec_znx_mut(),
 80                    i,
 81                    &sk.data.as_vec_znx(),
 82                    i,
 83                );
 84            });
 85        }
 86
 87        self.gglwe_key_compressed_encrypt_sk(&mut res.key, sk, &sk_out, seed_xa, source_xe, scratch_1);
 88
 89        res.p = p;
 90    }
 91}
 92
 93impl<DataSelf: DataMut> AutomorphismKeyCompressed<DataSelf> {
 94    #[allow(clippy::too_many_arguments)]
 95    pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
 96        &mut self,
 97        module: &Module<B>,
 98        p: i64,
 99        sk: &GLWESecret<DataSk>,
100        seed_xa: [u8; 32],
101        source_xe: &mut Source,
102        scratch: &mut Scratch<B>,
103    ) where
104        Module<B>: GGLWEAutomorphismKeyCompressedEncryptSk<B>,
105    {
106        module.gglwe_automorphism_key_compressed_encrypt_sk(self, p, sk, seed_xa, source_xe, scratch);
107    }
108}