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