1use poulpy_hal::{
  2    api::{
  3        ModuleN, ScratchAvailable, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolBytesOf, SvpPrepare, VecZnxAddInplace, VecZnxAddNormal,
  4        VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf, VecZnxFillUniform,
  5        VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace,
  6        VecZnxSwitchRing,
  7    },
  8    layouts::{Backend, DataMut, GaloisElement, Module, Scratch},
  9    source::Source,
 10};
 11
 12use crate::{
 13    ScratchTakeCore,
 14    layouts::{
 15        AutomorphismKey, AutomorphismKeyToMut, GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, GLWESwitchingKey,
 16    },
 17};
 18
 19impl AutomorphismKey<Vec<u8>> {
 20
 21    pub fn encrypt_sk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 22    where
 23        A: GGLWEInfos,
 24        M: GGLWEAutomorphismKeyEncryptSk<BE>
 25    {
 26        module.gglwe_automorphism_key_encrypt_sk_tmp_bytes(infos)
 27    }
 28
 29    pub fn encrypt_pk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 30    where
 31        A: GGLWEInfos,
 32        M: GGLWEAutomorphismKeyEncryptPk<BE>
 33    {
 34        module.gglwe_automorphism_key_encrypt_pk_tmp_bytes(infos)
 35    }
 36}
 37
 38impl<DM: DataMut> AutomorphismKey<DM>
 39where
 40    Self: AutomorphismKeyToMut,
 41{
 42    pub fn encrypt_sk<S, BE: Backend>(
 43        &mut self,
 44        module: &Module<BE>,
 45        p: i64,
 46        sk: &S,
 47        source_xa: &mut Source,
 48        source_xe: &mut Source,
 49        scratch: &mut Scratch<BE>,
 50    ) where
 51        S: GLWESecretToRef,
 52        Module<BE>: GGLWEAutomorphismKeyEncryptSk<BE>,
 53    {
 54        module.gglwe_automorphism_key_encrypt_sk(self, p, sk, source_xa, source_xe, scratch);
 55    }
 56}
 57
 58pub trait GGLWEAutomorphismKeyEncryptSk<BE: Backend> {
 59    fn gglwe_automorphism_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 60    where
 61        A: GGLWEInfos;
 62
 63    fn gglwe_automorphism_key_encrypt_sk<A, B>(
 64        &self,
 65        res: &mut A,
 66        p: i64,
 67        sk: &B,
 68        source_xa: &mut Source,
 69        source_xe: &mut Source,
 70        scratch: &mut Scratch<BE>,
 71    ) where
 72        A: AutomorphismKeyToMut,
 73        B: GLWESecretToRef;
 74}
 75
 76impl<BE: Backend> GGLWEAutomorphismKeyEncryptSk<BE> for Module<BE>
 77where
 78    Module<BE>: ModuleN 
 79        + VecZnxAddScalarInplace
 80        + VecZnxDftBytesOf
 81        + VecZnxBigNormalize<BE>
 82        + VecZnxDftApply<BE>
 83        + SvpApplyDftToDftInplace<BE>
 84        + VecZnxIdftApplyConsume<BE>
 85        + VecZnxNormalizeTmpBytes
 86        + VecZnxFillUniform
 87        + VecZnxSubInplace
 88        + VecZnxAddInplace
 89        + VecZnxNormalizeInplace<BE>
 90        + VecZnxAddNormal
 91        + VecZnxNormalize<BE>
 92        + VecZnxSub
 93        + SvpPrepare<BE>
 94        + VecZnxSwitchRing
 95        + SvpPPolBytesOf
 96        + VecZnxAutomorphism
 97        + SvpPPolAlloc<BE>
 98        + GaloisElement,
 99    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
100{
101
102    fn gglwe_automorphism_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
103    where
104        A: GGLWEInfos,
105    {
106        assert_eq!(
107            infos.rank_in(),
108            infos.rank_out(),
109            "rank_in != rank_out is not supported for GGLWEAutomorphismKey"
110        );
111        GLWESwitchingKey::encrypt_sk_tmp_bytes(self, infos) + GLWESecret::bytes_of_from_infos(self, &infos.glwe_layout())
112    }
113
114    fn gglwe_automorphism_key_encrypt_sk<A, B>(
115        &self,
116        res: &mut A,
117        p: i64,
118        sk: &B,
119        source_xa: &mut Source,
120        source_xe: &mut Source,
121        scratch: &mut Scratch<BE>,
122    ) where
123        A: AutomorphismKeyToMut,
124        B: GLWESecretToRef,
125    {
126        let res: &mut AutomorphismKey<&mut [u8]> = &mut res.to_mut();
127        let sk: &GLWESecret<&[u8]> = &sk.to_ref();
128
129        #[cfg(debug_assertions)]
130        {
131            use crate::layouts::{GLWEInfos, LWEInfos};
132
133            assert_eq!(res.n(), sk.n());
134            assert_eq!(res.rank_out(), res.rank_in());
135            assert_eq!(sk.rank(), res.rank_out());
136            assert!(
137                scratch.available() >= AutomorphismKey::encrypt_sk_tmp_bytes(self, res),
138                "scratch.available(): {} < AutomorphismKey::encrypt_sk_tmp_bytes: {:?}",
139                scratch.available(),
140                AutomorphismKey::encrypt_sk_tmp_bytes(self, res)
141            )
142        }
143
144        let (mut sk_out, scratch_1) = scratch.take_glwe_secret(self, sk.rank());
145
146        {
147            (0..res.rank_out().into()).for_each(|i| {
148                self.vec_znx_automorphism(
149                    self.galois_element_inv(p),
150                    &mut sk_out.data.as_vec_znx_mut(),
151                    i,
152                    &sk.data.as_vec_znx(),
153                    i,
154                );
155            });
156        }
157
158        res.key
159            .encrypt_sk(self, sk, &sk_out, source_xa, source_xe, scratch_1);
160
161        res.p = p;
162    }
163}
164
165
166pub trait GGLWEAutomorphismKeyEncryptPk<BE: Backend> {
167    fn gglwe_automorphism_key_encrypt_pk_tmp_bytes<A>(&self, infos: &A) -> usize
168    where
169        A: GGLWEInfos;
170
171}
172
173impl<BE: Backend> GGLWEAutomorphismKeyEncryptPk<BE> for Module<BE>
174where
175    Module<BE>: ModuleN 
176        + VecZnxAddScalarInplace
177        + VecZnxDftBytesOf
178        + VecZnxBigNormalize<BE>
179        + VecZnxDftApply<BE>
180        + SvpApplyDftToDftInplace<BE>
181        + VecZnxIdftApplyConsume<BE>
182        + VecZnxNormalizeTmpBytes
183        + VecZnxFillUniform
184        + VecZnxSubInplace
185        + VecZnxAddInplace
186        + VecZnxNormalizeInplace<BE>
187        + VecZnxAddNormal
188        + VecZnxNormalize<BE>
189        + VecZnxSub
190        + SvpPPolBytesOf
191        + SvpPPolAlloc<BE>,
192    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
193{
194
195    fn gglwe_automorphism_key_encrypt_pk_tmp_bytes<A>(&self, infos: &A) -> usize
196    where
197        A: GGLWEInfos,
198    {
199        assert_eq!(
200            infos.rank_in(),
201            infos.rank_out(),
202            "rank_in != rank_out is not supported for GGLWEAutomorphismKey"
203        );
204        GLWESwitchingKey::encrypt_pk_tmp_bytes(self, infos)
205    }
206}