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