1use poulpy_hal::{
  2    api::{
  3        ModuleN, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolBytesOf, SvpPrepare, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace,
  4        VecZnxAutomorphismInplace, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf, VecZnxFillUniform,
  5        VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace,
  6        VecZnxSwitchRing,
  7    },
  8    layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxView, ZnxViewMut},
  9    source::Source,
 10};
 11
 12use crate::{
 13    layouts::{
 14        GLWESecretAlloc,
 15        GGLWEInfos, GLWESecret, GLWESwitchingKey, LWEInfos, LWESecret, LWESwitchingKey, Rank,
 16        prepared::{GLWESecretPrepared, GLWESecretPreparedAlloc},
 17    },
 18    encryption::gglwe_ksk::GLWESwitchingKeyEncryptSk,
 19    ScratchTakeCore, 
 20};
 21
 22impl LWESwitchingKey<Vec<u8>> {
 23    pub fn encrypt_sk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 24    where
 25        A: GGLWEInfos,
 26        M: LWESwitchingKeyEncrypt<BE>,
 27    {
 28        module.lwe_switching_key_encrypt_sk_tmp_bytes(infos)
 29    }
 30    
 31}
 32
 33impl<D: DataMut> LWESwitchingKey<D> {
 34    #[allow(clippy::too_many_arguments)]
 35    pub fn encrypt_sk<DIn, DOut, B: Backend>(
 36        &mut self,
 37        module: &Module<B>,
 38        sk_lwe_in: &LWESecret<DIn>,
 39        sk_lwe_out: &LWESecret<DOut>,
 40        source_xa: &mut Source,
 41        source_xe: &mut Source,
 42        scratch: &mut Scratch<B>,
 43    ) where
 44        DIn: DataRef,
 45        DOut: DataRef,
 46        Module<B>: ModuleN
 47            + VecZnxAutomorphismInplace<B>
 48            + VecZnxAddScalarInplace
 49            + VecZnxDftBytesOf
 50            + VecZnxBigNormalize<B>
 51            + VecZnxDftApply<B>
 52            + SvpApplyDftToDftInplace<B>
 53            + VecZnxIdftApplyConsume<B>
 54            + VecZnxNormalizeTmpBytes
 55            + VecZnxFillUniform
 56            + VecZnxSubInplace
 57            + VecZnxAddInplace
 58            + VecZnxNormalizeInplace<B>
 59            + VecZnxAddNormal
 60            + VecZnxNormalize<B>
 61            + VecZnxSub
 62            + SvpPrepare<B>
 63            + VecZnxSwitchRing
 64            + SvpPPolBytesOf
 65            + SvpPPolAlloc<B>,
 66        Scratch<B>: ScratchTakeCore<B>,
 67    {
 68        #[cfg(debug_assertions)]
 69        {
 70            assert!(sk_lwe_in.n().0 <= self.n().0);
 71            assert!(sk_lwe_out.n().0 <= self.n().0);
 72            assert!(self.n().0 <= module.n() as u32);
 73        }
 74
 75        let (mut sk_in_glwe, scratch_1) = scratch.take_glwe_secret(module, Rank(1));
 76        let (mut sk_out_glwe, scratch_2) = scratch_1.take_glwe_secret(module, Rank(1));
 77
 78        sk_out_glwe.data.at_mut(0, 0)[..sk_lwe_out.n().into()].copy_from_slice(sk_lwe_out.data.at(0, 0));
 79        sk_out_glwe.data.at_mut(0, 0)[sk_lwe_out.n().into()..].fill(0);
 80        module.vec_znx_automorphism_inplace(-1, &mut sk_out_glwe.data.as_vec_znx_mut(), 0, scratch_2);
 81
 82        sk_in_glwe.data.at_mut(0, 0)[..sk_lwe_in.n().into()].copy_from_slice(sk_lwe_in.data.at(0, 0));
 83        sk_in_glwe.data.at_mut(0, 0)[sk_lwe_in.n().into()..].fill(0);
 84        module.vec_znx_automorphism_inplace(-1, &mut sk_in_glwe.data.as_vec_znx_mut(), 0, scratch_2);
 85
 86        self.0.encrypt_sk(
 87            module,
 88            &sk_in_glwe,
 89            &sk_out_glwe,
 90            source_xa,
 91            source_xe,
 92            scratch_2,
 93        );
 94    }
 95}
 96
 97pub trait LWESwitchingKeyEncrypt<BE: Backend>
 98where 
 99    Self: Sized
100        + ModuleN
101        + SvpPPolBytesOf
102        + SvpPPolAlloc<BE>
103        + VecZnxNormalizeTmpBytes
104        + VecZnxDftBytesOf
105        + VecZnxAutomorphismInplace<BE>
106        + VecZnxAddScalarInplace
107        + VecZnxBigNormalize<BE>
108        + VecZnxDftApply<BE>
109        + SvpApplyDftToDftInplace<BE>
110        + VecZnxIdftApplyConsume<BE>
111        + VecZnxFillUniform
112        + VecZnxSubInplace
113        + VecZnxAddInplace
114        + VecZnxNormalizeInplace<BE>
115        + VecZnxAddNormal
116        + VecZnxNormalize<BE>
117        + VecZnxSub
118        + SvpPrepare<BE>
119        + VecZnxSwitchRing
120        + GLWESecretAlloc
121        + GLWESecretPreparedAlloc<BE>
122{
123    fn lwe_switching_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
124    where
125        A: GGLWEInfos;
126}
127
128impl<BE:Backend> LWESwitchingKeyEncrypt<BE> for Module<BE> where 
129    Self: ModuleN
130        + SvpPPolBytesOf
131        + SvpPPolAlloc<BE>
132        + VecZnxNormalizeTmpBytes
133        + VecZnxDftBytesOf
134        + VecZnxAutomorphismInplace<BE>
135        + VecZnxAddScalarInplace
136        + VecZnxBigNormalize<BE>
137        + VecZnxDftApply<BE>
138        + SvpApplyDftToDftInplace<BE>
139        + VecZnxIdftApplyConsume<BE>
140        + VecZnxFillUniform
141        + VecZnxSubInplace
142        + VecZnxAddInplace
143        + VecZnxNormalizeInplace<BE>
144        + VecZnxAddNormal
145        + VecZnxNormalize<BE>
146        + VecZnxSub
147        + SvpPrepare<BE>
148        + VecZnxSwitchRing
149        + GLWESecretAlloc
150        + GLWESecretPreparedAlloc<BE>
151        + GLWESwitchingKeyEncryptSk<BE>
152        ,
153{
154    fn lwe_switching_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
155    where
156        A: GGLWEInfos
157    {
158        debug_assert_eq!(
159            infos.dsize().0,
160            1,
161            "dsize > 1 is not supported for LWESwitchingKey"
162        );
163        debug_assert_eq!(
164            infos.rank_in().0,
165            1,
166            "rank_in > 1 is not supported for LWESwitchingKey"
167        );
168        debug_assert_eq!(
169            infos.rank_out().0,
170            1,
171            "rank_out > 1 is not supported for LWESwitchingKey"
172        );
173        GLWESecret::bytes_of(self, Rank(1))
174            + GLWESecretPrepared::bytes_of(self, Rank(1))
175            + GLWESwitchingKey::encrypt_sk_tmp_bytes(self, infos)
176    }
177}