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