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