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        GLWESecretPreparedApi, 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: AutomorphismKeyCompressedEncryptSk<BE>,
 20    {
 21        module.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: AutomorphismKeyCompressedEncryptSk<BE>,
 38    {
 39        module.automorphism_key_compressed_encrypt_sk(self, p, sk, seed_xa, source_xe, scratch);
 40    }
 41}
 42
 43pub trait AutomorphismKeyCompressedEncryptSk<BE: Backend> {
 44    fn automorphism_key_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 45    where
 46        A: GGLWEInfos;
 47
 48    fn 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> AutomorphismKeyCompressedEncryptSk<BE> for Module<BE>
 62where
 63    Self: ModuleN + GaloisElement + VecZnxAutomorphism + GGLWECompressedEncryptSk<BE> + GLWESecretPreparedApi<BE>,
 64    Scratch<BE>: ScratchTakeCore<BE>,
 65{
 66    fn 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 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, sk.rank());
102            for i in 0..res.rank_out().into() {
103                self.vec_znx_automorphism(
104                    self.galois_element_inv(p),
105                    &mut sk_out.data.as_vec_znx_mut(),
106                    i,
107                    &sk.data.as_vec_znx(),
108                    i,
109                );
110            }
111            sk_out_prepared.prepare(self, &sk_out);
112        }
113
114        self.gglwe_compressed_encrypt_sk(
115            res,
116            &sk.data,
117            &sk_out_prepared,
118            seed_xa,
119            source_xe,
120            scratch_1,
121        );
122
123        res.set_p(p);
124    }
125}