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, GLWESecretPrepare, GLWESecretPrepared,
 11        GLWESecretPreparedAlloc, GLWESecretToRef, LWEInfos, SetAutomorphismGaloisElement,
 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 + SetAutomorphismGaloisElement + GGLWEInfos,
 68        S: GLWESecretToRef;
 69}
 70
 71impl<BE: Backend> AutomorphismKeyEncryptSk<BE> for Module<BE>
 72where
 73    Self: GGLWEEncryptSk<BE>
 74        + VecZnxAutomorphism
 75        + GaloisElement
 76        + SvpPPolBytesOf
 77        + GLWESecretPrepare<BE>
 78        + GLWESecretPreparedAlloc<BE>,
 79    Scratch<BE>: ScratchTakeCore<BE>,
 80{
 81    fn automorphism_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 82    where
 83        A: GGLWEInfos,
 84    {
 85        assert_eq!(
 86            infos.rank_in(),
 87            infos.rank_out(),
 88            "rank_in != rank_out is not supported for GGLWEAutomorphismKey"
 89        );
 90        GLWESecretPrepared::bytes_of_from_infos(self, infos)
 91            + self
 92                .gglwe_encrypt_sk_tmp_bytes(infos)
 93                .max(GLWESecret::bytes_of_from_infos(infos))
 94    }
 95
 96    fn automorphism_key_encrypt_sk<R, S>(
 97        &self,
 98        res: &mut R,
 99        p: i64,
100        sk: &S,
101        source_xa: &mut Source,
102        source_xe: &mut Source,
103        scratch: &mut Scratch<BE>,
104    ) where
105        R: GGLWEToMut + SetAutomorphismGaloisElement + GGLWEInfos,
106        S: GLWESecretToRef,
107    {
108        let sk: &GLWESecret<&[u8]> = &sk.to_ref();
109
110        assert_eq!(res.n(), sk.n());
111        assert_eq!(res.rank_out(), res.rank_in());
112        assert_eq!(sk.rank(), res.rank_out());
113        assert!(
114            scratch.available() >= self.automorphism_key_encrypt_sk_tmp_bytes(res),
115            "scratch.available(): {} < AutomorphismKey::encrypt_sk_tmp_bytes: {:?}",
116            scratch.available(),
117            self.automorphism_key_encrypt_sk_tmp_bytes(res)
118        );
119
120        let (mut sk_out_prepared, scratch_1) = scratch.take_glwe_secret_prepared(self, sk.rank());
121
122        {
123            let (mut sk_out, _) = scratch_1.take_glwe_secret(self, sk.rank());
124            (0..res.rank_out().into()).for_each(|i| {
125                self.vec_znx_automorphism(
126                    self.galois_element_inv(p),
127                    &mut sk_out.data.as_vec_znx_mut(),
128                    i,
129                    &sk.data.as_vec_znx(),
130                    i,
131                );
132            });
133            sk_out_prepared.prepare(self, &sk_out);
134        }
135
136        self.gglwe_encrypt_sk(
137            res,
138            &sk.data,
139            &sk_out_prepared,
140            source_xa,
141            source_xe,
142            scratch_1,
143        );
144
145        res.set_p(p);
146    }
147}
148
149pub trait GGLWEAutomorphismKeyEncryptPk<BE: Backend> {
150    fn automorphism_key_encrypt_pk_tmp_bytes<A>(&self, infos: &A) -> usize
151    where
152        A: GGLWEInfos;
153}
154
155impl<BE: Backend> GGLWEAutomorphismKeyEncryptPk<BE> for Module<BE>
156where
157    Self:,
158    Scratch<BE>: ScratchTakeCore<BE>,
159{
160    fn automorphism_key_encrypt_pk_tmp_bytes<A>(&self, _infos: &A) -> usize
161    where
162        A: GGLWEInfos,
163    {
164        unimplemented!()
165    }
166}