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