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