1use poulpy_hal::{
  2    api::{ScratchAvailable, SvpPPolBytesOf, VecZnxAutomorphism},
  3    layouts::{Backend, DataMut, GaloisElement, Module, Scratch},
  4    source::Source,
  5};
  6
  7use crate::{
  8    GGLWEEncryptSk, ScratchTakeCore,
  9    layouts::{
 10        GGLWEInfos, GGLWEToMut, GGLWEToRef, GLWEAutomorphismKey, GLWEInfos, GLWESecret, GLWESecretPrepared,
 11        GLWESecretPreparedFactory, GLWESecretToRef, LWEInfos, SetGaloisElement,
 12    },
 13};
 14
 15impl GLWEAutomorphismKey<Vec<u8>> {
 16    pub fn encrypt_sk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 17    where
 18        A: GGLWEInfos,
 19        M: GLWEAutomorphismKeyEncryptSk<BE>,
 20    {
 21        module.glwe_automorphism_key_encrypt_sk_tmp_bytes(infos)
 22    }
 23
 24    pub fn encrypt_pk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 25    where
 26        A: GGLWEInfos,
 27        M: GLWEAutomorphismKeyEncryptPk<BE>,
 28    {
 29        module.glwe_automorphism_key_encrypt_pk_tmp_bytes(infos)
 30    }
 31}
 32
 33impl<DM: DataMut> GLWEAutomorphismKey<DM>
 34where
 35    Self: GGLWEToRef,
 36{
 37    pub fn encrypt_sk<S, M, BE: Backend>(
 38        &mut self,
 39        module: &M,
 40        p: i64,
 41        sk: &S,
 42        source_xa: &mut Source,
 43        source_xe: &mut Source,
 44        scratch: &mut Scratch<BE>,
 45    ) where
 46        S: GLWESecretToRef,
 47        M: GLWEAutomorphismKeyEncryptSk<BE>,
 48    {
 49        module.glwe_automorphism_key_encrypt_sk(self, p, sk, source_xa, source_xe, scratch);
 50    }
 51}
 52
 53pub trait GLWEAutomorphismKeyEncryptSk<BE: Backend> {
 54    fn glwe_automorphism_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 55    where
 56        A: GGLWEInfos;
 57
 58    fn glwe_automorphism_key_encrypt_sk<R, S>(
 59        &self,
 60        res: &mut R,
 61        p: i64,
 62        sk: &S,
 63        source_xa: &mut Source,
 64        source_xe: &mut Source,
 65        scratch: &mut Scratch<BE>,
 66    ) where
 67        R: GGLWEToMut + SetGaloisElement + GGLWEInfos,
 68        S: GLWESecretToRef;
 69}
 70
 71impl<BE: Backend> GLWEAutomorphismKeyEncryptSk<BE> for Module<BE>
 72where
 73    Self: GGLWEEncryptSk<BE> + VecZnxAutomorphism + GaloisElement + SvpPPolBytesOf + GLWESecretPreparedFactory<BE>,
 74    Scratch<BE>: ScratchTakeCore<BE>,
 75{
 76    fn glwe_automorphism_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 77    where
 78        A: GGLWEInfos,
 79    {
 80        assert_eq!(
 81            infos.rank_in(),
 82            infos.rank_out(),
 83            "rank_in != rank_out is not supported for GGLWEAutomorphismKey"
 84        );
 85        GLWESecretPrepared::bytes_of_from_infos(self, infos)
 86            + self
 87                .gglwe_encrypt_sk_tmp_bytes(infos)
 88                .max(GLWESecret::bytes_of_from_infos(infos))
 89    }
 90
 91    fn glwe_automorphism_key_encrypt_sk<R, S>(
 92        &self,
 93        res: &mut R,
 94        p: i64,
 95        sk: &S,
 96        source_xa: &mut Source,
 97        source_xe: &mut Source,
 98        scratch: &mut Scratch<BE>,
 99    ) where
100        R: GGLWEToMut + SetGaloisElement + GGLWEInfos,
101        S: GLWESecretToRef,
102    {
103        let sk: &GLWESecret<&[u8]> = &sk.to_ref();
104
105        assert_eq!(res.n(), sk.n());
106        assert_eq!(res.rank_out(), res.rank_in());
107        assert_eq!(sk.rank(), res.rank_out());
108        assert!(
109            scratch.available() >= self.glwe_automorphism_key_encrypt_sk_tmp_bytes(res),
110            "scratch.available(): {} < AutomorphismKey::encrypt_sk_tmp_bytes: {:?}",
111            scratch.available(),
112            self.glwe_automorphism_key_encrypt_sk_tmp_bytes(res)
113        );
114
115        let (mut sk_out_prepared, scratch_1) = scratch.take_glwe_secret_prepared(self, sk.rank());
116
117        {
118            let (mut sk_out, _) = scratch_1.take_glwe_secret(sk.n(), sk.rank());
119            sk_out.dist = sk.dist;
120
121            for i in 0..sk.rank().into() {
122                self.vec_znx_automorphism(
123                    self.galois_element_inv(p),
124                    &mut sk_out.data.as_vec_znx_mut(),
125                    i,
126                    &sk.data.as_vec_znx(),
127                    i,
128                );
129            }
130            sk_out_prepared.prepare(self, &sk_out);
131        }
132
133        self.gglwe_encrypt_sk(
134            res,
135            &sk.data,
136            &sk_out_prepared,
137            source_xa,
138            source_xe,
139            scratch_1,
140        );
141
142        res.set_p(p);
143    }
144}
145
146pub trait GLWEAutomorphismKeyEncryptPk<BE: Backend> {
147    fn glwe_automorphism_key_encrypt_pk_tmp_bytes<A>(&self, infos: &A) -> usize
148    where
149        A: GGLWEInfos;
150}
151
152impl<BE: Backend> GLWEAutomorphismKeyEncryptPk<BE> for Module<BE>
153where
154    Self:,
155    Scratch<BE>: ScratchTakeCore<BE>,
156{
157    fn glwe_automorphism_key_encrypt_pk_tmp_bytes<A>(&self, _infos: &A) -> usize
158    where
159        A: GGLWEInfos,
160    {
161        unimplemented!()
162    }
163}