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