1use backend::{Encoding, FFT64, Module, ScratchOwned, ZnxView};
  2use sampling::source::Source;
  3
  4use crate::{
  5    FourierGLWESecret, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos, LWECiphertext, LWESecret,
  6    lwe::{
  7        LWEPlaintext,
  8        keyswtich::{GLWEToLWESwitchingKey, LWESwitchingKey, LWEToGLWESwitchingKey},
  9    },
 10};
 11
 12#[test]
 13fn lwe_to_glwe() {
 14    let n: usize = 1 << 5;
 15    let module: Module<FFT64> = Module::<FFT64>::new(n);
 16    let basek: usize = 17;
 17    let sigma: f64 = 3.2;
 18
 19    let rank: usize = 2;
 20
 21    let n_lwe: usize = 22;
 22    let k_lwe_ct: usize = 2 * basek;
 23    let k_lwe_pt: usize = 8;
 24
 25    let k_glwe_ct: usize = 3 * basek;
 26
 27    let k_ksk: usize = k_lwe_ct + basek;
 28
 29    let mut source_xs: Source = Source::new([0u8; 32]);
 30    let mut source_xa: Source = Source::new([0u8; 32]);
 31    let mut source_xe: Source = Source::new([0u8; 32]);
 32
 33    let mut scratch: ScratchOwned = ScratchOwned::new(
 34        LWEToGLWESwitchingKey::encrypt_sk_scratch_space(&module, basek, k_ksk, rank)
 35            | GLWECiphertext::from_lwe_scratch_space(&module, basek, k_lwe_ct, k_glwe_ct, k_ksk, rank)
 36            | GLWECiphertext::decrypt_scratch_space(&module, basek, k_glwe_ct),
 37    );
 38
 39    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
 40    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
 41
 42    let mut sk_glwe_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::alloc(&module, rank);
 43    sk_glwe_dft.set(&module, &sk_glwe);
 44
 45    let mut sk_lwe = LWESecret::alloc(n_lwe);
 46    sk_lwe.fill_ternary_prob(0.5, &mut source_xs);
 47
 48    let data: i64 = 17;
 49
 50    let mut lwe_pt: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_pt);
 51    lwe_pt
 52        .data
 53        .encode_coeff_i64(0, basek, k_lwe_pt, 0, data, k_lwe_pt);
 54
 55    let mut lwe_ct: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe_ct);
 56    lwe_ct.encrypt_sk(&lwe_pt, &sk_lwe, &mut source_xa, &mut source_xe, sigma);
 57
 58    let mut ksk: LWEToGLWESwitchingKey<Vec<u8>, FFT64> = LWEToGLWESwitchingKey::alloc(&module, basek, k_ksk, lwe_ct.size(), rank);
 59
 60    ksk.encrypt_sk(
 61        &module,
 62        &sk_lwe,
 63        &sk_glwe_dft,
 64        &mut source_xa,
 65        &mut source_xe,
 66        sigma,
 67        scratch.borrow(),
 68    );
 69
 70    let mut glwe_ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_glwe_ct, rank);
 71    glwe_ct.from_lwe(&module, &lwe_ct, &ksk, scratch.borrow());
 72
 73    let mut glwe_pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_glwe_ct);
 74    glwe_ct.decrypt(&module, &mut glwe_pt, &sk_glwe_dft, scratch.borrow());
 75
 76    assert_eq!(glwe_pt.data.at(0, 0)[0], lwe_pt.data.at(0, 0)[0]);
 77}
 78
 79#[test]
 80fn glwe_to_lwe() {
 81    let n: usize = 1 << 5;
 82    let module: Module<FFT64> = Module::<FFT64>::new(n);
 83    let basek: usize = 17;
 84    let sigma: f64 = 3.2;
 85
 86    let rank: usize = 2;
 87
 88    let n_lwe: usize = 22;
 89    let k_lwe_ct: usize = 2 * basek;
 90    let k_lwe_pt: usize = 8;
 91
 92    let k_glwe_ct: usize = 3 * basek;
 93
 94    let k_ksk: usize = k_lwe_ct + basek;
 95
 96    let mut source_xs: Source = Source::new([0u8; 32]);
 97    let mut source_xa: Source = Source::new([0u8; 32]);
 98    let mut source_xe: Source = Source::new([0u8; 32]);
 99
100    let mut scratch: ScratchOwned = ScratchOwned::new(
101        LWEToGLWESwitchingKey::encrypt_sk_scratch_space(&module, basek, k_ksk, rank)
102            | GLWECiphertext::from_lwe_scratch_space(&module, basek, k_lwe_ct, k_glwe_ct, k_ksk, rank)
103            | GLWECiphertext::decrypt_scratch_space(&module, basek, k_glwe_ct),
104    );
105
106    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
107    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
108
109    let mut sk_glwe_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::alloc(&module, rank);
110    sk_glwe_dft.set(&module, &sk_glwe);
111
112    let mut sk_lwe = LWESecret::alloc(n_lwe);
113    sk_lwe.fill_ternary_prob(0.5, &mut source_xs);
114
115    let data: i64 = 17;
116    let mut glwe_pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_glwe_ct);
117
118    glwe_pt
119        .data
120        .encode_coeff_i64(0, basek, k_lwe_pt, 0, data, k_lwe_pt);
121
122    let mut glwe_ct = GLWECiphertext::alloc(&module, basek, k_glwe_ct, rank);
123    glwe_ct.encrypt_sk(
124        &module,
125        &glwe_pt,
126        &sk_glwe_dft,
127        &mut source_xa,
128        &mut source_xe,
129        sigma,
130        scratch.borrow(),
131    );
132
133    let mut ksk: GLWEToLWESwitchingKey<Vec<u8>, FFT64> =
134        GLWEToLWESwitchingKey::alloc(&module, basek, k_ksk, glwe_ct.size(), rank);
135
136    ksk.encrypt_sk(
137        &module,
138        &sk_lwe,
139        &sk_glwe,
140        &mut source_xa,
141        &mut source_xe,
142        sigma,
143        scratch.borrow(),
144    );
145
146    let mut lwe_ct: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe_ct);
147    lwe_ct.from_glwe(&module, &glwe_ct, &ksk, scratch.borrow());
148
149    let mut lwe_pt: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_ct);
150    lwe_ct.decrypt(&mut lwe_pt, &sk_lwe);
151
152    assert_eq!(glwe_pt.data.at(0, 0)[0], lwe_pt.data.at(0, 0)[0]);
153}
154
155#[test]
156fn keyswitch() {
157    let n: usize = 1 << 5;
158    let module: Module<FFT64> = Module::<FFT64>::new(n);
159    let basek: usize = 17;
160    let sigma: f64 = 3.2;
161
162    let n_lwe_in: usize = 22;
163    let n_lwe_out: usize = 30;
164    let k_lwe_ct: usize = 2 * basek;
165    let k_lwe_pt: usize = 8;
166
167    let k_ksk: usize = k_lwe_ct + basek;
168
169    let mut source_xs: Source = Source::new([0u8; 32]);
170    let mut source_xa: Source = Source::new([0u8; 32]);
171    let mut source_xe: Source = Source::new([0u8; 32]);
172
173    let mut scratch: ScratchOwned = ScratchOwned::new(
174        LWESwitchingKey::encrypt_sk_scratch_space(&module, basek, k_ksk)
175            | LWECiphertext::keyswitch_scratch_space(&module, basek, k_lwe_ct, k_lwe_ct, k_ksk),
176    );
177
178    let mut sk_lwe_in: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_in);
179    sk_lwe_in.fill_ternary_prob(0.5, &mut source_xs);
180
181    let mut sk_lwe_out: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_out);
182    sk_lwe_out.fill_ternary_prob(0.5, &mut source_xs);
183
184    let data: i64 = 17;
185
186    let mut lwe_pt_in: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_pt);
187    lwe_pt_in
188        .data
189        .encode_coeff_i64(0, basek, k_lwe_pt, 0, data, k_lwe_pt);
190
191    let mut lwe_ct_in: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe_in, basek, k_lwe_ct);
192    lwe_ct_in.encrypt_sk(
193        &lwe_pt_in,
194        &sk_lwe_in,
195        &mut source_xa,
196        &mut source_xe,
197        sigma,
198    );
199
200    let mut ksk: LWESwitchingKey<Vec<u8>, FFT64> = LWESwitchingKey::alloc(&module, basek, k_ksk, lwe_ct_in.size());
201
202    ksk.encrypt_sk(
203        &module,
204        &sk_lwe_in,
205        &sk_lwe_out,
206        &mut source_xa,
207        &mut source_xe,
208        sigma,
209        scratch.borrow(),
210    );
211
212    let mut lwe_ct_out: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe_out, basek, k_lwe_ct);
213
214    lwe_ct_out.keyswitch(&module, &lwe_ct_in, &ksk, scratch.borrow());
215
216    let mut lwe_pt_out: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_ct);
217    lwe_ct_out.decrypt(&mut lwe_pt_out, &sk_lwe_out);
218
219    assert_eq!(lwe_pt_in.data.at(0, 0)[0], lwe_pt_out.data.at(0, 0)[0]);
220}