1use backend::{
  2    hal::{
  3        api::{
  4            MatZnxAlloc, ModuleNew, ScalarZnxAlloc, ScalarZnxAllocBytes, ScratchOwnedAlloc, ScratchOwnedBorrow,
  5            VecZnxAddScalarInplace, VecZnxAlloc, VecZnxAllocBytes, VecZnxAutomorphismInplace, VecZnxEncodeCoeffsi64,
  6            VecZnxSwithcDegree, ZnxView,
  7        },
  8        layouts::{Backend, Module, ScratchOwned},
  9        oep::{
 10            ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
 11            TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
 12        },
 13    },
 14    implementation::cpu_spqlios::FFT64,
 15};
 16use sampling::source::Source;
 17
 18use crate::{
 19    GGLWEEncryptSkFamily, GGLWEExecLayoutFamily, GLWECiphertext, GLWEDecryptFamily, GLWEKeyswitchFamily, GLWEPlaintext,
 20    GLWESecret, GLWESecretExec, Infos, LWECiphertext, LWESecret,
 21    lwe::{
 22        LWEPlaintext,
 23        keyswtich::{
 24            GLWEToLWESwitchingKey, GLWEToLWESwitchingKeyExec, LWESwitchingKey, LWESwitchingKeyExec, LWEToGLWESwitchingKey,
 25            LWEToGLWESwitchingKeyExec,
 26        },
 27    },
 28};
 29
 30#[test]
 31fn lwe_to_glwe() {
 32    let log_n: usize = 5;
 33    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
 34    test_lwe_to_glwe(&module)
 35}
 36
 37pub(crate) trait LWETestModuleFamily<B: Backend> = GGLWEEncryptSkFamily<B>
 38    + GLWEDecryptFamily<B>
 39    + VecZnxSwithcDegree
 40    + VecZnxAddScalarInplace
 41    + VecZnxAlloc
 42    + GGLWEExecLayoutFamily<B>
 43    + GLWEKeyswitchFamily<B>
 44    + ScalarZnxAllocBytes
 45    + VecZnxAllocBytes
 46    + ScalarZnxAlloc
 47    + VecZnxEncodeCoeffsi64
 48    + MatZnxAlloc
 49    + VecZnxAutomorphismInplace;
 50
 51pub(crate) trait LWETestScratchFamily<B: Backend> = TakeScalarZnxImpl<B>
 52    + TakeVecZnxDftImpl<B>
 53    + ScratchAvailableImpl<B>
 54    + TakeVecZnxImpl<B>
 55    + TakeVecZnxBigImpl<B>
 56    + TakeSvpPPolImpl<B>
 57    + ScratchOwnedAllocImpl<B>
 58    + ScratchOwnedBorrowImpl<B>;
 59
 60pub(crate) fn test_lwe_to_glwe<B: Backend>(module: &Module<B>)
 61where
 62    Module<B>: LWETestModuleFamily<B>,
 63    B: LWETestScratchFamily<B>,
 64{
 65    let basek: usize = 17;
 66    let sigma: f64 = 3.2;
 67
 68    let rank: usize = 2;
 69
 70    let n_lwe: usize = 22;
 71    let k_lwe_ct: usize = 2 * basek;
 72    let k_lwe_pt: usize = 8;
 73
 74    let k_glwe_ct: usize = 3 * basek;
 75
 76    let k_ksk: usize = k_lwe_ct + basek;
 77
 78    let mut source_xs: Source = Source::new([0u8; 32]);
 79    let mut source_xa: Source = Source::new([0u8; 32]);
 80    let mut source_xe: Source = Source::new([0u8; 32]);
 81
 82    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 83        LWEToGLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k_ksk, rank)
 84            | GLWECiphertext::from_lwe_scratch_space(module, basek, k_lwe_ct, k_glwe_ct, k_ksk, rank)
 85            | GLWECiphertext::decrypt_scratch_space(module, basek, k_glwe_ct),
 86    );
 87
 88    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank);
 89    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
 90
 91    let sk_glwe_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_glwe);
 92
 93    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
 94    sk_lwe.fill_ternary_prob(0.5, &mut source_xs);
 95
 96    let data: i64 = 17;
 97
 98    let mut lwe_pt: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_pt);
 99    module.encode_coeff_i64(basek, &mut lwe_pt.data, 0, k_lwe_pt, 0, data, k_lwe_pt);
100
101    let mut lwe_ct: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe_ct);
102    lwe_ct.encrypt_sk(
103        module,
104        &lwe_pt,
105        &sk_lwe,
106        &mut source_xa,
107        &mut source_xe,
108        sigma,
109    );
110
111    let mut ksk: LWEToGLWESwitchingKey<Vec<u8>> = LWEToGLWESwitchingKey::alloc(module, basek, k_ksk, lwe_ct.size(), rank);
112
113    ksk.encrypt_sk(
114        module,
115        &sk_lwe,
116        &sk_glwe,
117        &mut source_xa,
118        &mut source_xe,
119        sigma,
120        scratch.borrow(),
121    );
122
123    let mut glwe_ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_glwe_ct, rank);
124
125    let ksk_exec: LWEToGLWESwitchingKeyExec<Vec<u8>, B> = LWEToGLWESwitchingKeyExec::from(module, &ksk, scratch.borrow());
126
127    glwe_ct.from_lwe(module, &lwe_ct, &ksk_exec, scratch.borrow());
128
129    let mut glwe_pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_glwe_ct);
130    glwe_ct.decrypt(module, &mut glwe_pt, &sk_glwe_exec, scratch.borrow());
131
132    assert_eq!(glwe_pt.data.at(0, 0)[0], lwe_pt.data.at(0, 0)[0]);
133}
134
135#[test]
136fn glwe_to_lwe() {
137    let log_n: usize = 5;
138    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
139    test_glwe_to_lwe(&module)
140}
141
142fn test_glwe_to_lwe<B: Backend>(module: &Module<B>)
143where
144    Module<B>: LWETestModuleFamily<B>,
145    B: LWETestScratchFamily<B>,
146{
147    let basek: usize = 17;
148    let sigma: f64 = 3.2;
149
150    let rank: usize = 2;
151
152    let n_lwe: usize = 22;
153    let k_lwe_ct: usize = 2 * basek;
154    let k_lwe_pt: usize = 8;
155
156    let k_glwe_ct: usize = 3 * basek;
157
158    let k_ksk: usize = k_lwe_ct + basek;
159
160    let mut source_xs: Source = Source::new([0u8; 32]);
161    let mut source_xa: Source = Source::new([0u8; 32]);
162    let mut source_xe: Source = Source::new([0u8; 32]);
163
164    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
165        LWEToGLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k_ksk, rank)
166            | LWECiphertext::from_glwe_scratch_space(module, basek, k_lwe_ct, k_glwe_ct, k_ksk, rank)
167            | GLWECiphertext::decrypt_scratch_space(module, basek, k_glwe_ct),
168    );
169
170    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank);
171    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
172
173    let sk_glwe_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_glwe);
174
175    let mut sk_lwe = LWESecret::alloc(n_lwe);
176    sk_lwe.fill_ternary_prob(0.5, &mut source_xs);
177
178    let data: i64 = 17;
179    let mut glwe_pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_glwe_ct);
180    module.encode_coeff_i64(basek, &mut glwe_pt.data, 0, k_lwe_pt, 0, data, k_lwe_pt);
181
182    let mut glwe_ct = GLWECiphertext::alloc(module, basek, k_glwe_ct, rank);
183    glwe_ct.encrypt_sk(
184        module,
185        &glwe_pt,
186        &sk_glwe_exec,
187        &mut source_xa,
188        &mut source_xe,
189        sigma,
190        scratch.borrow(),
191    );
192
193    let mut ksk: GLWEToLWESwitchingKey<Vec<u8>> = GLWEToLWESwitchingKey::alloc(module, basek, k_ksk, glwe_ct.size(), rank);
194
195    ksk.encrypt_sk(
196        module,
197        &sk_lwe,
198        &sk_glwe,
199        &mut source_xa,
200        &mut source_xe,
201        sigma,
202        scratch.borrow(),
203    );
204
205    let mut lwe_ct: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe_ct);
206
207    let ksk_exec: GLWEToLWESwitchingKeyExec<Vec<u8>, B> = GLWEToLWESwitchingKeyExec::from(module, &ksk, scratch.borrow());
208
209    lwe_ct.from_glwe(module, &glwe_ct, &ksk_exec, scratch.borrow());
210
211    let mut lwe_pt: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_ct);
212    lwe_ct.decrypt(module, &mut lwe_pt, &sk_lwe);
213
214    assert_eq!(glwe_pt.data.at(0, 0)[0], lwe_pt.data.at(0, 0)[0]);
215}
216
217#[test]
218fn keyswitch() {
219    let log_n: usize = 5;
220    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
221    test_keyswitch(&module)
222}
223
224fn test_keyswitch<B: Backend>(module: &Module<B>)
225where
226    Module<B>: LWETestModuleFamily<B>,
227    B: LWETestScratchFamily<B>,
228{
229    let basek: usize = 17;
230    let sigma: f64 = 3.2;
231
232    let n_lwe_in: usize = 22;
233    let n_lwe_out: usize = 30;
234    let k_lwe_ct: usize = 2 * basek;
235    let k_lwe_pt: usize = 8;
236
237    let k_ksk: usize = k_lwe_ct + basek;
238
239    let mut source_xs: Source = Source::new([0u8; 32]);
240    let mut source_xa: Source = Source::new([0u8; 32]);
241    let mut source_xe: Source = Source::new([0u8; 32]);
242
243    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
244        LWESwitchingKey::encrypt_sk_scratch_space(module, basek, k_ksk)
245            | LWECiphertext::keyswitch_scratch_space(module, basek, k_lwe_ct, k_lwe_ct, k_ksk),
246    );
247
248    let mut sk_lwe_in: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_in);
249    sk_lwe_in.fill_ternary_prob(0.5, &mut source_xs);
250
251    let mut sk_lwe_out: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_out);
252    sk_lwe_out.fill_ternary_prob(0.5, &mut source_xs);
253
254    let data: i64 = 17;
255
256    let mut lwe_pt_in: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_pt);
257    module.encode_coeff_i64(basek, &mut lwe_pt_in.data, 0, k_lwe_pt, 0, data, k_lwe_pt);
258
259    let mut lwe_ct_in: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe_in, basek, k_lwe_ct);
260    lwe_ct_in.encrypt_sk(
261        module,
262        &lwe_pt_in,
263        &sk_lwe_in,
264        &mut source_xa,
265        &mut source_xe,
266        sigma,
267    );
268
269    let mut ksk: LWESwitchingKey<Vec<u8>> = LWESwitchingKey::alloc(module, basek, k_ksk, lwe_ct_in.size());
270
271    ksk.encrypt_sk(
272        module,
273        &sk_lwe_in,
274        &sk_lwe_out,
275        &mut source_xa,
276        &mut source_xe,
277        sigma,
278        scratch.borrow(),
279    );
280
281    let mut lwe_ct_out: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe_out, basek, k_lwe_ct);
282
283    let ksk_exec: LWESwitchingKeyExec<Vec<u8>, B> = LWESwitchingKeyExec::from(module, &ksk, scratch.borrow());
284
285    lwe_ct_out.keyswitch(module, &lwe_ct_in, &ksk_exec, scratch.borrow());
286
287    let mut lwe_pt_out: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_ct);
288    lwe_ct_out.decrypt(module, &mut lwe_pt_out, &sk_lwe_out);
289
290    assert_eq!(lwe_pt_in.data.at(0, 0)[0], lwe_pt_out.data.at(0, 0)[0]);
291}