1use poulpy_hal::{
  2    api::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow},
  3    layouts::{Backend, Module, Scratch, ScratchOwned, ZnxView},
  4    source::Source,
  5};
  6
  7use crate::{
  8    GLWEDecrypt, GLWEEncryptSk, GLWEFromLWE, GLWEToLWESwitchingKeyEncryptSk, LWEDecrypt, LWEEncryptSk, LWEFromGLWE,
  9    LWEToGLWESwitchingKeyEncryptSk, ScratchTakeCore,
 10    layouts::{
 11        Base2K, Degree, Dnum, GLWE, GLWELayout, GLWEPlaintext, GLWESecret, GLWESecretPreparedFactory, GLWEToLWEKey,
 12        GLWEToLWEKeyLayout, GLWEToLWEKeyPrepared, GLWEToLWEKeyPreparedFactory, LWE, LWELayout, LWEPlaintext, LWESecret,
 13        LWEToGLWEKey, LWEToGLWEKeyLayout, LWEToGLWEKeyPrepared, LWEToGLWEKeyPreparedFactory, Rank, TorusPrecision,
 14        prepared::GLWESecretPrepared,
 15    },
 16};
 17
 18pub fn test_lwe_to_glwe<BE: Backend>(module: &Module<BE>)
 19where
 20    Module<BE>: GLWEFromLWE<BE>
 21        + LWEToGLWESwitchingKeyEncryptSk<BE>
 22        + GLWEDecrypt<BE>
 23        + GLWESecretPreparedFactory<BE>
 24        + LWEEncryptSk<BE>
 25        + LWEToGLWEKeyPreparedFactory<BE>,
 26    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 27    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 28{
 29    let n_glwe: Degree = Degree(module.n() as u32);
 30    let n_lwe: Degree = Degree(22);
 31
 32    let rank: Rank = Rank(2);
 33    let k_lwe_pt: TorusPrecision = TorusPrecision(8);
 34
 35    let mut source_xs: Source = Source::new([0u8; 32]);
 36    let mut source_xa: Source = Source::new([0u8; 32]);
 37    let mut source_xe: Source = Source::new([0u8; 32]);
 38
 39    let lwe_to_glwe_infos: LWEToGLWEKeyLayout = LWEToGLWEKeyLayout {
 40        n: n_glwe,
 41        base2k: Base2K(17),
 42        k: TorusPrecision(51),
 43        dnum: Dnum(2),
 44        rank_out: rank,
 45    };
 46
 47    let glwe_infos: GLWELayout = GLWELayout {
 48        n: n_glwe,
 49        base2k: Base2K(17),
 50        k: TorusPrecision(34),
 51        rank,
 52    };
 53
 54    let lwe_infos: LWELayout = LWELayout {
 55        n: n_lwe,
 56        base2k: Base2K(17),
 57        k: TorusPrecision(34),
 58    };
 59
 60    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
 61        LWEToGLWEKey::encrypt_sk_tmp_bytes(module, &lwe_to_glwe_infos)
 62            | GLWE::from_lwe_tmp_bytes(module, &glwe_infos, &lwe_infos, &lwe_to_glwe_infos)
 63            | GLWE::decrypt_tmp_bytes(module, &glwe_infos),
 64    );
 65
 66    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_infos);
 67    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
 68
 69    let mut sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc_from_infos(module, &sk_glwe);
 70    sk_glwe_prepared.prepare(module, &sk_glwe);
 71
 72    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
 73    sk_lwe.fill_ternary_prob(0.5, &mut source_xs);
 74
 75    let data: i64 = 17;
 76
 77    let mut lwe_pt: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc_from_infos(&lwe_infos);
 78    lwe_pt.encode_i64(data, k_lwe_pt);
 79
 80    let mut lwe_ct: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_infos);
 81    lwe_ct.encrypt_sk(module, &lwe_pt, &sk_lwe, &mut source_xa, &mut source_xe);
 82
 83    let mut ksk: LWEToGLWEKey<Vec<u8>> = LWEToGLWEKey::alloc_from_infos(&lwe_to_glwe_infos);
 84
 85    ksk.encrypt_sk(
 86        module,
 87        &sk_lwe,
 88        &sk_glwe_prepared,
 89        &mut source_xa,
 90        &mut source_xe,
 91        scratch.borrow(),
 92    );
 93
 94    let mut glwe_ct: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_infos);
 95
 96    let mut ksk_prepared: LWEToGLWEKeyPrepared<Vec<u8>, BE> = LWEToGLWEKeyPrepared::alloc_from_infos(module, &ksk);
 97    ksk_prepared.prepare(module, &ksk, scratch.borrow());
 98
 99    glwe_ct.from_lwe(module, &lwe_ct, &ksk_prepared, scratch.borrow());
100
101    let mut glwe_pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_infos);
102    glwe_ct.decrypt(module, &mut glwe_pt, &sk_glwe_prepared, scratch.borrow());
103
104    assert_eq!(glwe_pt.data.at(0, 0)[0], lwe_pt.data.at(0, 0)[0]);
105}
106
107pub fn test_glwe_to_lwe<BE: Backend>(module: &Module<BE>)
108where
109    Module<BE>: GLWEFromLWE<BE>
110        + GLWEToLWESwitchingKeyEncryptSk<BE>
111        + GLWEEncryptSk<BE>
112        + LWEDecrypt<BE>
113        + LWEFromGLWE<BE>
114        + GLWEDecrypt<BE>
115        + GLWESecretPreparedFactory<BE>
116        + GLWEToLWESwitchingKeyEncryptSk<BE>
117        + GLWEToLWEKeyPreparedFactory<BE>,
118    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
119    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
120{
121    let n_glwe: Degree = Degree(module.n() as u32);
122    let n_lwe: Degree = Degree(22);
123
124    let rank: Rank = Rank(2);
125    let k_lwe_pt: TorusPrecision = TorusPrecision(8);
126
127    let glwe_to_lwe_infos: GLWEToLWEKeyLayout = GLWEToLWEKeyLayout {
128        n: n_glwe,
129        base2k: Base2K(17),
130        k: TorusPrecision(51),
131        dnum: Dnum(2),
132        rank_in: rank,
133    };
134
135    let glwe_infos: GLWELayout = GLWELayout {
136        n: n_glwe,
137        base2k: Base2K(17),
138        k: TorusPrecision(34),
139        rank,
140    };
141
142    let lwe_infos: LWELayout = LWELayout {
143        n: n_lwe,
144        base2k: Base2K(17),
145        k: TorusPrecision(34),
146    };
147
148    let mut source_xs: Source = Source::new([0u8; 32]);
149    let mut source_xa: Source = Source::new([0u8; 32]);
150    let mut source_xe: Source = Source::new([0u8; 32]);
151
152    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
153        GLWEToLWEKey::encrypt_sk_tmp_bytes(module, &glwe_to_lwe_infos)
154            | LWE::from_glwe_tmp_bytes(module, &lwe_infos, &glwe_infos, &glwe_to_lwe_infos)
155            | GLWE::decrypt_tmp_bytes(module, &glwe_infos),
156    );
157
158    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_infos);
159    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
160
161    let mut sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc_from_infos(module, &sk_glwe);
162    sk_glwe_prepared.prepare(module, &sk_glwe);
163
164    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
165    sk_lwe.fill_ternary_prob(0.5, &mut source_xs);
166
167    let a_idx: usize = 1;
168
169    let mut data: Vec<i64> = vec![0i64; module.n()];
170    data[a_idx] = 17;
171    let mut glwe_pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_infos);
172    glwe_pt.encode_vec_i64(&data, k_lwe_pt);
173
174    println!("glwe_pt: {glwe_pt}");
175
176    let mut glwe_ct: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_infos);
177    glwe_ct.encrypt_sk(
178        module,
179        &glwe_pt,
180        &sk_glwe_prepared,
181        &mut source_xa,
182        &mut source_xe,
183        scratch.borrow(),
184    );
185
186    let mut ksk: GLWEToLWEKey<Vec<u8>> = GLWEToLWEKey::alloc_from_infos(&glwe_to_lwe_infos);
187
188    ksk.encrypt_sk(
189        module,
190        &sk_lwe,
191        &sk_glwe,
192        &mut source_xa,
193        &mut source_xe,
194        scratch.borrow(),
195    );
196
197    let mut lwe_ct: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_infos);
198
199    let mut ksk_prepared: GLWEToLWEKeyPrepared<Vec<u8>, BE> = GLWEToLWEKeyPrepared::alloc_from_infos(module, &ksk);
200    ksk_prepared.prepare(module, &ksk, scratch.borrow());
201
202    lwe_ct.from_glwe(module, &glwe_ct, a_idx, &ksk_prepared, scratch.borrow());
203
204    let mut lwe_pt: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc_from_infos(&lwe_infos);
205    lwe_ct.decrypt(module, &mut lwe_pt, &sk_lwe);
206
207    assert_eq!(glwe_pt.data.at(0, 0)[a_idx], lwe_pt.data.at(0, 0)[0]);
208}