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        AutomorphismKey, GGLWEInfos, GGLWEToMut, GGLWEToRef, GLWEInfos, GLWESecret, GLWESecretPrepared,
 11        GLWESecretPreparedFactory, GLWESecretToRef, LWEInfos, SetGaloisElement,
 12    },
 13};
 14
 15impl AutomorphismKey<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: AutomorphismKeyEncryptSk<BE>,
 20    {
 21        module.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: GGLWEAutomorphismKeyEncryptPk<BE>,
 28    {
 29        module.automorphism_key_encrypt_pk_tmp_bytes(infos)
 30    }
 31}
 32
 33impl<DM: DataMut> AutomorphismKey<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: AutomorphismKeyEncryptSk<BE>,
 48    {
 49        module.automorphism_key_encrypt_sk(self, p, sk, source_xa, source_xe, scratch);
 50    }
 51}
 52
 53pub trait AutomorphismKeyEncryptSk<BE: Backend> {
 54    fn automorphism_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 55    where
 56        A: GGLWEInfos;
 57
 58    fn 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> AutomorphismKeyEncryptSk<BE> for Module<BE>
 72where
 73    Self: GGLWEEncryptSk<BE> + VecZnxAutomorphism + GaloisElement + SvpPPolBytesOf + GLWESecretPreparedFactory<BE>,
 74    Scratch<BE>: ScratchTakeCore<BE>,
 75{
 76    fn 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 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.automorphism_key_encrypt_sk_tmp_bytes(res),
110            "scratch.available(): {} < AutomorphismKey::encrypt_sk_tmp_bytes: {:?}",
111            scratch.available(),
112            self.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            for i in 0..sk.rank().into() {
120                self.vec_znx_automorphism(
121                    self.galois_element_inv(p),
122                    &mut sk_out.data.as_vec_znx_mut(),
123                    i,
124                    &sk.data.as_vec_znx(),
125                    i,
126                );
127            }
128            sk_out_prepared.prepare(self, &sk_out);
129        }
130
131        self.gglwe_encrypt_sk(
132            res,
133            &sk.data,
134            &sk_out_prepared,
135            source_xa,
136            source_xe,
137            scratch_1,
138        );
139
140        res.set_p(p);
141    }
142}
143
144pub trait GGLWEAutomorphismKeyEncryptPk<BE: Backend> {
145    fn automorphism_key_encrypt_pk_tmp_bytes<A>(&self, infos: &A) -> usize
146    where
147        A: GGLWEInfos;
148}
149
150impl<BE: Backend> GGLWEAutomorphismKeyEncryptPk<BE> for Module<BE>
151where
152    Self:,
153    Scratch<BE>: ScratchTakeCore<BE>,
154{
155    fn automorphism_key_encrypt_pk_tmp_bytes<A>(&self, _infos: &A) -> usize
156    where
157        A: GGLWEInfos,
158    {
159        unimplemented!()
160    }
161}