1use poulpy_hal::{
  2    api::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxFillUniform, VecZnxNormalize},
  3    layouts::{Backend, FillUniform, Module, Scratch, ScratchOwned, ZnxView},
  4    source::Source,
  5};
  6use rug::Float;
  7
  8use crate::{
  9    GLWEDecrypt, GLWEEncryptSk, GLWEFromLWE, GLWENoise, GLWENormalize, GLWEToLWESwitchingKeyEncryptSk, LWEDecrypt, LWEEncryptSk,
 10    LWEFromGLWE, LWEToGLWESwitchingKeyEncryptSk, SIGMA, ScratchTakeCore,
 11    layouts::{
 12        Base2K, Degree, Dnum, GLWE, GLWELayout, GLWEPlaintext, GLWESecret, GLWESecretPreparedFactory, GLWEToLWEKey,
 13        GLWEToLWEKeyLayout, GLWEToLWEKeyPrepared, GLWEToLWEKeyPreparedFactory, LWE, LWEInfos, LWELayout, LWEPlaintext, LWESecret,
 14        LWEToGLWEKey, LWEToGLWEKeyLayout, LWEToGLWEKeyPrepared, LWEToGLWEKeyPreparedFactory, Rank, TorusPrecision,
 15        prepared::GLWESecretPrepared,
 16    },
 17};
 18
 19pub fn test_glwe_base2k_conversion<BE: Backend>(module: &Module<BE>)
 20where
 21    Module<BE>: GLWEEncryptSk<BE>
 22        + GLWEDecrypt<BE>
 23        + GLWENormalize<BE>
 24        + VecZnxFillUniform
 25        + GLWESecretPreparedFactory<BE>
 26        + GLWENoise<BE>,
 27    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 28    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 29{
 30    let n_glwe: Degree = Degree(module.n() as u32);
 31
 32    let mut source_xs: Source = Source::new([0u8; 32]);
 33    let mut source_xa: Source = Source::new([0u8; 32]);
 34    let mut source_xe: Source = Source::new([0u8; 32]);
 35
 36    for rank in 1_usize..3 {
 37        for bases in [[12, 8], [8, 12]] {
 38            let glwe_infos_in: GLWELayout = GLWELayout {
 39                n: n_glwe,
 40                base2k: Base2K(bases[0]),
 41                k: TorusPrecision(34),
 42                rank: Rank(rank as u32),
 43            };
 44
 45            let glwe_infos_out: GLWELayout = GLWELayout {
 46                n: n_glwe,
 47                base2k: Base2K(bases[1]),
 48                k: TorusPrecision(34),
 49                rank: Rank(rank as u32),
 50            };
 51
 52            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(module.n().into(), rank.into());
 53            sk.fill_ternary_prob(0.5, &mut source_xs);
 54
 55            let mut sk_prep: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc_from_infos(module, &sk);
 56            sk_prep.prepare(module, &sk);
 57
 58            let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
 59                GLWE::encrypt_sk_tmp_bytes(module, &glwe_infos_in).max(GLWE::decrypt_tmp_bytes(module, &glwe_infos_out)),
 60            );
 61
 62            let mut ct_in: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_infos_in);
 63            let mut ct_out: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_infos_out);
 64
 65            let pt_in: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_infos_in);
 66            let pt_out: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_infos_in);
 67
 68            ct_in.encrypt_sk(
 69                module,
 70                &pt_in,
 71                &sk_prep,
 72                &mut source_xa,
 73                &mut source_xe,
 74                scratch.borrow(),
 75            );
 76
 77            let mut data: Vec<Float> = (0..module.n()).map(|_| Float::with_val(128, 0)).collect();
 78            ct_in
 79                .data()
 80                .decode_vec_float(ct_in.base2k().into(), 0, &mut data);
 81
 82            ct_out.fill_uniform(ct_out.base2k().into(), &mut source_xa);
 83            module.glwe_normalize(&mut ct_out, &ct_in, scratch.borrow());
 84
 85            let mut data_conv: Vec<Float> = (0..module.n()).map(|_| Float::with_val(128, 0)).collect();
 86            ct_out
 87                .data()
 88                .decode_vec_float(ct_out.base2k().into(), 0, &mut data_conv);
 89
 90            ct_out.assert_noise(
 91                module,
 92                &sk_prep,
 93                &pt_out,
 94                -(ct_out.k().as_u32() as f64) + SIGMA.log2() + 0.5,
 95            );
 96        }
 97    }
 98}
 99
100pub fn test_lwe_to_glwe<BE: Backend>(module: &Module<BE>)
101where
102    Module<BE>: GLWEFromLWE<BE>
103        + LWEToGLWESwitchingKeyEncryptSk<BE>
104        + GLWEDecrypt<BE>
105        + GLWESecretPreparedFactory<BE>
106        + LWEEncryptSk<BE>
107        + LWEToGLWEKeyPreparedFactory<BE>
108        + VecZnxNormalize<BE>,
109    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
110    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
111{
112    let n_glwe: Degree = Degree(module.n() as u32);
113    let n_lwe: Degree = Degree(22);
114
115    let rank: Rank = Rank(2);
116    let k_lwe_pt: TorusPrecision = TorusPrecision(8);
117
118    let mut source_xs: Source = Source::new([0u8; 32]);
119    let mut source_xa: Source = Source::new([0u8; 32]);
120    let mut source_xe: Source = Source::new([0u8; 32]);
121
122    let lwe_to_glwe_infos: LWEToGLWEKeyLayout = LWEToGLWEKeyLayout {
123        n: n_glwe,
124        base2k: Base2K(13),
125        k: TorusPrecision(92),
126        dnum: Dnum(2),
127        rank_out: rank,
128    };
129
130    let glwe_infos: GLWELayout = GLWELayout {
131        n: n_glwe,
132        base2k: Base2K(15),
133        k: TorusPrecision(75),
134        rank,
135    };
136
137    let lwe_infos: LWELayout = LWELayout {
138        n: n_lwe,
139        base2k: Base2K(17),
140        k: TorusPrecision(75),
141    };
142
143    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
144        LWEToGLWEKey::encrypt_sk_tmp_bytes(module, &lwe_to_glwe_infos)
145            | GLWE::from_lwe_tmp_bytes(module, &glwe_infos, &lwe_infos, &lwe_to_glwe_infos)
146            | GLWE::decrypt_tmp_bytes(module, &glwe_infos),
147    );
148
149    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_infos);
150    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
151
152    let mut sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc_from_infos(module, &sk_glwe);
153    sk_glwe_prepared.prepare(module, &sk_glwe);
154
155    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
156    sk_lwe.fill_ternary_prob(0.5, &mut source_xs);
157
158    let data: i64 = 17;
159
160    let mut lwe_pt: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc_from_infos(&lwe_infos);
161    lwe_pt.encode_i64(data, k_lwe_pt);
162
163    let mut lwe_ct: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_infos);
164    lwe_ct.encrypt_sk(
165        module,
166        &lwe_pt,
167        &sk_lwe,
168        &mut source_xa,
169        &mut source_xe,
170        scratch.borrow(),
171    );
172
173    let mut ksk: LWEToGLWEKey<Vec<u8>> = LWEToGLWEKey::alloc_from_infos(&lwe_to_glwe_infos);
174
175    ksk.encrypt_sk(
176        module,
177        &sk_lwe,
178        &sk_glwe_prepared,
179        &mut source_xa,
180        &mut source_xe,
181        scratch.borrow(),
182    );
183
184    let mut glwe_ct: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_infos);
185
186    let mut ksk_prepared: LWEToGLWEKeyPrepared<Vec<u8>, BE> = LWEToGLWEKeyPrepared::alloc_from_infos(module, &ksk);
187    ksk_prepared.prepare(module, &ksk, scratch.borrow());
188
189    glwe_ct.from_lwe(module, &lwe_ct, &ksk_prepared, scratch.borrow());
190
191    let mut glwe_pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_infos);
192    glwe_ct.decrypt(module, &mut glwe_pt, &sk_glwe_prepared, scratch.borrow());
193
194    let mut lwe_pt_conv = LWEPlaintext::alloc(glwe_pt.base2k(), lwe_pt.k());
195
196    module.vec_znx_normalize(
197        glwe_pt.base2k().as_usize(),
198        lwe_pt_conv.data_mut(),
199        0,
200        lwe_pt.base2k().as_usize(),
201        lwe_pt.data(),
202        0,
203        scratch.borrow(),
204    );
205
206    assert_eq!(glwe_pt.data.at(0, 0)[0], lwe_pt_conv.data.at(0, 0)[0]);
207}
208
209pub fn test_glwe_to_lwe<BE: Backend>(module: &Module<BE>)
210where
211    Module<BE>: GLWEFromLWE<BE>
212        + GLWEToLWESwitchingKeyEncryptSk<BE>
213        + GLWEEncryptSk<BE>
214        + LWEDecrypt<BE>
215        + LWEFromGLWE<BE>
216        + GLWEDecrypt<BE>
217        + GLWESecretPreparedFactory<BE>
218        + GLWEToLWESwitchingKeyEncryptSk<BE>
219        + GLWEToLWEKeyPreparedFactory<BE>
220        + VecZnxNormalize<BE>,
221    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
222    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
223{
224    let n_glwe: Degree = Degree(module.n() as u32);
225    let n_lwe: Degree = Degree(22);
226
227    let rank: Rank = Rank(2);
228    let k_lwe_pt: TorusPrecision = TorusPrecision(8);
229
230    let glwe_to_lwe_infos: GLWEToLWEKeyLayout = GLWEToLWEKeyLayout {
231        n: n_glwe,
232        base2k: Base2K(13),
233        k: TorusPrecision(91),
234        dnum: Dnum(2),
235        rank_in: rank,
236    };
237
238    let glwe_infos: GLWELayout = GLWELayout {
239        n: n_glwe,
240        base2k: Base2K(17),
241        k: TorusPrecision(72),
242        rank,
243    };
244
245    let lwe_infos: LWELayout = LWELayout {
246        n: n_lwe,
247        base2k: Base2K(15),
248        k: TorusPrecision(72),
249    };
250
251    let mut source_xs: Source = Source::new([0u8; 32]);
252    let mut source_xa: Source = Source::new([0u8; 32]);
253    let mut source_xe: Source = Source::new([0u8; 32]);
254
255    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
256        GLWEToLWEKey::encrypt_sk_tmp_bytes(module, &glwe_to_lwe_infos)
257            | LWE::from_glwe_tmp_bytes(module, &lwe_infos, &glwe_infos, &glwe_to_lwe_infos)
258            | GLWE::decrypt_tmp_bytes(module, &glwe_infos),
259    );
260
261    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_infos);
262    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
263
264    let mut sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc_from_infos(module, &sk_glwe);
265    sk_glwe_prepared.prepare(module, &sk_glwe);
266
267    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
268    sk_lwe.fill_ternary_prob(0.5, &mut source_xs);
269
270    let a_idx: usize = 1;
271
272    let mut data: Vec<i64> = vec![0i64; module.n()];
273    data[a_idx] = 17;
274    let mut glwe_pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_infos);
275    glwe_pt.encode_vec_i64(&data, k_lwe_pt);
276
277    println!("glwe_pt: {glwe_pt}");
278
279    let mut glwe_ct: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_infos);
280    glwe_ct.encrypt_sk(
281        module,
282        &glwe_pt,
283        &sk_glwe_prepared,
284        &mut source_xa,
285        &mut source_xe,
286        scratch.borrow(),
287    );
288
289    let mut ksk: GLWEToLWEKey<Vec<u8>> = GLWEToLWEKey::alloc_from_infos(&glwe_to_lwe_infos);
290
291    ksk.encrypt_sk(
292        module,
293        &sk_lwe,
294        &sk_glwe,
295        &mut source_xa,
296        &mut source_xe,
297        scratch.borrow(),
298    );
299
300    let mut lwe_ct: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_infos);
301
302    let mut ksk_prepared: GLWEToLWEKeyPrepared<Vec<u8>, BE> = GLWEToLWEKeyPrepared::alloc_from_infos(module, &ksk);
303    ksk_prepared.prepare(module, &ksk, scratch.borrow());
304
305    lwe_ct.from_glwe(module, &glwe_ct, a_idx, &ksk_prepared, scratch.borrow());
306
307    let mut lwe_pt: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc_from_infos(&lwe_infos);
308    lwe_ct.decrypt(module, &mut lwe_pt, &sk_lwe, scratch.borrow());
309
310    let mut glwe_pt_conv = GLWEPlaintext::alloc(glwe_ct.n(), lwe_pt.base2k(), lwe_pt.k());
311
312    module.vec_znx_normalize(
313        lwe_pt.base2k().as_usize(),
314        glwe_pt_conv.data_mut(),
315        0,
316        glwe_ct.base2k().as_usize(),
317        glwe_pt.data(),
318        0,
319        scratch.borrow(),
320    );
321
322    assert_eq!(glwe_pt_conv.data.at(0, 0)[a_idx], lwe_pt.data.at(0, 0)[0]);
323}