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