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