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, GLWESecretPrepare, GLWESecretPrepared,
 11        GLWESecretPreparedAlloc, GLWESecretToRef, LWEInfos, SetAutomorphismGaloisElement, compressed::AutomorphismKeyCompressed,
 12    },
 13};
 14
 15impl AutomorphismKeyCompressed<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> AutomorphismKeyCompressed<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 + SetAutomorphismGaloisElement + GGLWEInfos,
 58        S: GLWESecretToRef + GLWEInfos;
 59}
 60
 61impl<BE: Backend> AutomorphismKeyCompressedEncryptSk<BE> for Module<BE>
 62where
 63    Self: ModuleN
 64        + GaloisElement
 65        + VecZnxAutomorphism
 66        + GGLWECompressedEncryptSk<BE>
 67        + GLWESecretPreparedAlloc<BE>
 68        + GLWESecretPrepare<BE>,
 69    Scratch<BE>: ScratchTakeCore<BE>,
 70{
 71    fn automorphism_key_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 72    where
 73        A: GGLWEInfos,
 74    {
 75        assert_eq!(self.n() as u32, infos.n());
 76        self.gglwe_compressed_encrypt_sk_tmp_bytes(infos)
 77            .max(GLWESecret::bytes_of_from_infos(infos))
 78            + GLWESecretPrepared::bytes_of_from_infos(self, infos)
 79    }
 80
 81    fn automorphism_key_compressed_encrypt_sk<R, S>(
 82        &self,
 83        res: &mut R,
 84        p: i64,
 85        sk: &S,
 86        seed_xa: [u8; 32],
 87        source_xe: &mut Source,
 88        scratch: &mut Scratch<BE>,
 89    ) where
 90        R: GGLWECompressedToMut + GGLWECompressedSeedMut + SetAutomorphismGaloisElement + GGLWEInfos,
 91        S: GLWESecretToRef + GLWEInfos,
 92    {
 93        let sk: &GLWESecret<&[u8]> = &sk.to_ref();
 94        assert_eq!(res.n(), sk.n());
 95        assert_eq!(res.rank_out(), res.rank_in());
 96        assert_eq!(sk.rank(), res.rank_out());
 97        assert!(
 98            scratch.available() >= AutomorphismKeyCompressed::encrypt_sk_tmp_bytes(self, res),
 99            "scratch.available(): {} < AutomorphismKey::encrypt_sk_tmp_bytes: {}",
100            scratch.available(),
101            AutomorphismKeyCompressed::encrypt_sk_tmp_bytes(self, res)
102        );
103
104        let (mut sk_out_prepared, scratch_1) = scratch.take_glwe_secret_prepared(self, sk.rank());
105        {
106            let (mut sk_out, _) = scratch_1.take_glwe_secret(self, sk.rank());
107            for i in 0..res.rank_out().into() {
108                self.vec_znx_automorphism(
109                    self.galois_element_inv(p),
110                    &mut sk_out.data.as_vec_znx_mut(),
111                    i,
112                    &sk.data.as_vec_znx(),
113                    i,
114                );
115            }
116            sk_out_prepared.prepare(self, &sk_out);
117        }
118
119        self.gglwe_compressed_encrypt_sk(
120            res,
121            &sk.data,
122            &sk_out_prepared,
123            seed_xa,
124            source_xe,
125            scratch_1,
126        );
127
128        res.set_p(p);
129    }
130}