1use backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxAutomorphismInplace, VecZnxSwithcDegree,
  4        VmpPMatAlloc, VmpPMatPrepare, ZnxView,
  5    },
  6    layouts::{Backend, Module, ScratchOwned},
  7    oep::{
  8        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
  9        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
 10    },
 11};
 12use sampling::source::Source;
 13
 14use crate::layouts::{
 15    GLWECiphertext, GLWEPlaintext, GLWESecret, GLWEToLWESwitchingKey, Infos, LWECiphertext, LWEPlaintext, LWESecret,
 16    LWEToGLWESwitchingKey,
 17    prepared::{GLWESecretPrepared, GLWEToLWESwitchingKeyPrepared, LWEToGLWESwitchingKeyPrepared, PrepareAlloc},
 18};
 19
 20use crate::trait_families::{GGLWEEncryptSkFamily, GLWEDecryptFamily, GLWEKeyswitchFamily, GLWESecretPreparedModuleFamily};
 21
 22pub fn test_lwe_to_glwe<B: Backend>(module: &Module<B>)
 23where
 24    Module<B>: GGLWEEncryptSkFamily<B>
 25        + GLWEDecryptFamily<B>
 26        + VecZnxSwithcDegree
 27        + VecZnxAddScalarInplace
 28        + GLWEKeyswitchFamily<B>
 29        + VecZnxAutomorphismInplace
 30        + VmpPMatAlloc<B>
 31        + VmpPMatPrepare<B>
 32        + GLWESecretPreparedModuleFamily<B>,
 33    B: TakeScalarZnxImpl<B>
 34        + TakeVecZnxDftImpl<B>
 35        + ScratchAvailableImpl<B>
 36        + TakeVecZnxImpl<B>
 37        + TakeVecZnxBigImpl<B>
 38        + TakeSvpPPolImpl<B>
 39        + ScratchOwnedAllocImpl<B>
 40        + ScratchOwnedBorrowImpl<B>,
 41{
 42    let n: usize = module.n();
 43    let basek: usize = 17;
 44    let sigma: f64 = 3.2;
 45
 46    let rank: usize = 2;
 47
 48    let n_lwe: usize = 22;
 49    let k_lwe_ct: usize = 2 * basek;
 50    let k_lwe_pt: usize = 8;
 51
 52    let k_glwe_ct: usize = 3 * basek;
 53
 54    let k_ksk: usize = k_lwe_ct + basek;
 55
 56    let mut source_xs: Source = Source::new([0u8; 32]);
 57    let mut source_xa: Source = Source::new([0u8; 32]);
 58    let mut source_xe: Source = Source::new([0u8; 32]);
 59
 60    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 61        LWEToGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank)
 62            | GLWECiphertext::from_lwe_scratch_space(module, n, basek, k_lwe_ct, k_glwe_ct, k_ksk, rank)
 63            | GLWECiphertext::decrypt_scratch_space(module, n, basek, k_glwe_ct),
 64    );
 65
 66    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
 67    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
 68
 69    let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch.borrow());
 70
 71    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
 72    sk_lwe.fill_ternary_prob(0.5, &mut source_xs);
 73
 74    let data: i64 = 17;
 75
 76    let mut lwe_pt: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_pt);
 77    lwe_pt.encode_i64(data, k_lwe_pt);
 78
 79    let mut lwe_ct: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe_ct);
 80    lwe_ct.encrypt_sk(
 81        module,
 82        &lwe_pt,
 83        &sk_lwe,
 84        &mut source_xa,
 85        &mut source_xe,
 86        sigma,
 87    );
 88
 89    let mut ksk: LWEToGLWESwitchingKey<Vec<u8>> = LWEToGLWESwitchingKey::alloc(n, basek, k_ksk, lwe_ct.size(), rank);
 90
 91    ksk.encrypt_sk(
 92        module,
 93        &sk_lwe,
 94        &sk_glwe,
 95        &mut source_xa,
 96        &mut source_xe,
 97        sigma,
 98        scratch.borrow(),
 99    );
100
101    let mut glwe_ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_glwe_ct, rank);
102
103    let ksk_prepared: LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
104
105    glwe_ct.from_lwe(module, &lwe_ct, &ksk_prepared, scratch.borrow());
106
107    let mut glwe_pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_glwe_ct);
108    glwe_ct.decrypt(module, &mut glwe_pt, &sk_glwe_prepared, scratch.borrow());
109
110    assert_eq!(glwe_pt.data.at(0, 0)[0], lwe_pt.data.at(0, 0)[0]);
111}
112
113pub fn test_glwe_to_lwe<B: Backend>(module: &Module<B>)
114where
115    Module<B>: GGLWEEncryptSkFamily<B>
116        + GLWEDecryptFamily<B>
117        + VecZnxSwithcDegree
118        + VecZnxAddScalarInplace
119        + GLWEKeyswitchFamily<B>
120        + VecZnxAutomorphismInplace
121        + VmpPMatAlloc<B>
122        + VmpPMatPrepare<B>
123        + GLWESecretPreparedModuleFamily<B>,
124    B: TakeScalarZnxImpl<B>
125        + TakeVecZnxDftImpl<B>
126        + ScratchAvailableImpl<B>
127        + TakeVecZnxImpl<B>
128        + TakeVecZnxBigImpl<B>
129        + TakeSvpPPolImpl<B>
130        + ScratchOwnedAllocImpl<B>
131        + ScratchOwnedBorrowImpl<B>,
132{
133    let n: usize = module.n();
134    let basek: usize = 17;
135    let sigma: f64 = 3.2;
136
137    let rank: usize = 2;
138
139    let n_lwe: usize = 22;
140    let k_lwe_ct: usize = 2 * basek;
141    let k_lwe_pt: usize = 8;
142
143    let k_glwe_ct: usize = 3 * basek;
144
145    let k_ksk: usize = k_lwe_ct + basek;
146
147    let mut source_xs: Source = Source::new([0u8; 32]);
148    let mut source_xa: Source = Source::new([0u8; 32]);
149    let mut source_xe: Source = Source::new([0u8; 32]);
150
151    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
152        LWEToGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank)
153            | LWECiphertext::from_glwe_scratch_space(module, n, basek, k_lwe_ct, k_glwe_ct, k_ksk, rank)
154            | GLWECiphertext::decrypt_scratch_space(module, n, basek, k_glwe_ct),
155    );
156
157    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
158    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
159
160    let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch.borrow());
161
162    let mut sk_lwe = LWESecret::alloc(n_lwe);
163    sk_lwe.fill_ternary_prob(0.5, &mut source_xs);
164
165    let data: i64 = 17;
166    let mut glwe_pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_glwe_ct);
167    glwe_pt.encode_coeff_i64(data, k_lwe_pt, 0);
168
169    let mut glwe_ct = GLWECiphertext::alloc(n, basek, k_glwe_ct, rank);
170    glwe_ct.encrypt_sk(
171        module,
172        &glwe_pt,
173        &sk_glwe_prepared,
174        &mut source_xa,
175        &mut source_xe,
176        sigma,
177        scratch.borrow(),
178    );
179
180    let mut ksk: GLWEToLWESwitchingKey<Vec<u8>> = GLWEToLWESwitchingKey::alloc(n, basek, k_ksk, glwe_ct.size(), rank);
181
182    ksk.encrypt_sk(
183        module,
184        &sk_lwe,
185        &sk_glwe,
186        &mut source_xa,
187        &mut source_xe,
188        sigma,
189        scratch.borrow(),
190    );
191
192    let mut lwe_ct: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe_ct);
193
194    let ksk_prepared: GLWEToLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
195
196    lwe_ct.from_glwe(module, &glwe_ct, &ksk_prepared, scratch.borrow());
197
198    let mut lwe_pt: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_ct);
199    lwe_ct.decrypt(module, &mut lwe_pt, &sk_lwe);
200
201    assert_eq!(glwe_pt.data.at(0, 0)[0], lwe_pt.data.at(0, 0)[0]);
202}