1use backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxFillUniform, VecZnxSwithcDegree, VmpPMatAlloc,
  4        VmpPMatPrepare,
  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::{
 15    layouts::{
 16        GGLWESwitchingKey, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos,
 17        prepared::{GGLWESwitchingKeyPrepared, GLWESecretPrepared, PrepareAlloc},
 18    },
 19    noise::log2_std_noise_gglwe_product,
 20    trait_families::{GLWEDecryptFamily, GLWEKeyswitchFamily},
 21};
 22
 23use crate::trait_families::{GGLWESwitchingKeyEncryptSkFamily, GLWESecretPreparedModuleFamily};
 24
 25pub fn test_glwe_keyswitch<B: Backend>(
 26    module: &Module<B>,
 27    basek: usize,
 28    k_out: usize,
 29    k_in: usize,
 30    k_ksk: usize,
 31    digits: usize,
 32    rank_in: usize,
 33    rank_out: usize,
 34    sigma: f64,
 35) where
 36    Module<B>: GGLWESwitchingKeyEncryptSkFamily<B>
 37        + GLWESecretPreparedModuleFamily<B>
 38        + GLWEKeyswitchFamily<B>
 39        + GLWEDecryptFamily<B>
 40        + VecZnxSwithcDegree
 41        + VecZnxAddScalarInplace
 42        + VmpPMatAlloc<B>
 43        + VmpPMatPrepare<B>,
 44    B: TakeVecZnxDftImpl<B>
 45        + TakeVecZnxBigImpl<B>
 46        + TakeSvpPPolImpl<B>
 47        + ScratchOwnedAllocImpl<B>
 48        + ScratchOwnedBorrowImpl<B>
 49        + ScratchAvailableImpl<B>
 50        + TakeScalarZnxImpl<B>
 51        + TakeVecZnxImpl<B>,
 52{
 53    let n = module.n();
 54    let rows: usize = k_in.div_ceil(basek * digits);
 55
 56    let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank_in, rank_out);
 57    let mut ct_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_in, rank_in);
 58    let mut ct_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_out, rank_out);
 59    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_in);
 60
 61    let mut source_xs: Source = Source::new([0u8; 32]);
 62    let mut source_xe: Source = Source::new([0u8; 32]);
 63    let mut source_xa: Source = Source::new([0u8; 32]);
 64
 65    module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_in, &mut source_xa);
 66
 67    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 68        GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, ksk.k(), rank_in, rank_out)
 69            | GLWECiphertext::encrypt_sk_scratch_space(module, n, basek, ct_in.k())
 70            | GLWECiphertext::keyswitch_scratch_space(
 71                module,
 72                n,
 73                basek,
 74                ct_out.k(),
 75                ct_in.k(),
 76                ksk.k(),
 77                digits,
 78                rank_in,
 79                rank_out,
 80            ),
 81    );
 82
 83    let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
 84    sk_in.fill_ternary_prob(0.5, &mut source_xs);
 85    let sk_in_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
 86
 87    let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
 88    sk_out.fill_ternary_prob(0.5, &mut source_xs);
 89    let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
 90
 91    ksk.encrypt_sk(
 92        module,
 93        &sk_in,
 94        &sk_out,
 95        &mut source_xa,
 96        &mut source_xe,
 97        sigma,
 98        scratch.borrow(),
 99    );
100
101    ct_in.encrypt_sk(
102        module,
103        &pt_want,
104        &sk_in_prepared,
105        &mut source_xa,
106        &mut source_xe,
107        sigma,
108        scratch.borrow(),
109    );
110
111    let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
112
113    ct_out.keyswitch(module, &ct_in, &ksk_prepared, scratch.borrow());
114
115    let max_noise: f64 = log2_std_noise_gglwe_product(
116        module.n() as f64,
117        basek * digits,
118        0.5,
119        0.5,
120        0f64,
121        sigma * sigma,
122        0f64,
123        rank_in as f64,
124        k_in,
125        k_ksk,
126    );
127
128    ct_out.assert_noise(module, &sk_out_prepared, &pt_want, max_noise + 0.5);
129}
130
131pub fn test_glwe_keyswitch_inplace<B: Backend>(
132    module: &Module<B>,
133    basek: usize,
134    k_ct: usize,
135    k_ksk: usize,
136    digits: usize,
137    rank: usize,
138    sigma: f64,
139) where
140    Module<B>: GGLWESwitchingKeyEncryptSkFamily<B>
141        + GLWESecretPreparedModuleFamily<B>
142        + GLWEKeyswitchFamily<B>
143        + GLWEDecryptFamily<B>
144        + VecZnxSwithcDegree
145        + VecZnxAddScalarInplace
146        + VmpPMatAlloc<B>
147        + VmpPMatPrepare<B>,
148    B: TakeVecZnxDftImpl<B>
149        + TakeVecZnxBigImpl<B>
150        + TakeSvpPPolImpl<B>
151        + ScratchOwnedAllocImpl<B>
152        + ScratchOwnedBorrowImpl<B>
153        + ScratchAvailableImpl<B>
154        + TakeScalarZnxImpl<B>
155        + TakeVecZnxImpl<B>,
156{
157    let n: usize = module.n();
158    let rows: usize = k_ct.div_ceil(basek * digits);
159
160    let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank, rank);
161    let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_ct, rank);
162    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_ct);
163
164    let mut source_xs: Source = Source::new([0u8; 32]);
165    let mut source_xe: Source = Source::new([0u8; 32]);
166    let mut source_xa: Source = Source::new([0u8; 32]);
167
168    module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_ct, &mut source_xa);
169
170    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
171        GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, ksk.k(), rank, rank)
172            | GLWECiphertext::encrypt_sk_scratch_space(module, n, basek, ct_glwe.k())
173            | GLWECiphertext::keyswitch_inplace_scratch_space(module, n, basek, ct_glwe.k(), ksk.k(), digits, rank),
174    );
175
176    let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
177    sk_in.fill_ternary_prob(0.5, &mut source_xs);
178    let sk_in_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
179
180    let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
181    sk_out.fill_ternary_prob(0.5, &mut source_xs);
182    let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
183
184    ksk.encrypt_sk(
185        module,
186        &sk_in,
187        &sk_out,
188        &mut source_xa,
189        &mut source_xe,
190        sigma,
191        scratch.borrow(),
192    );
193
194    ct_glwe.encrypt_sk(
195        module,
196        &pt_want,
197        &sk_in_prepared,
198        &mut source_xa,
199        &mut source_xe,
200        sigma,
201        scratch.borrow(),
202    );
203
204    let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
205
206    ct_glwe.keyswitch_inplace(module, &ksk_prepared, scratch.borrow());
207
208    let max_noise: f64 = log2_std_noise_gglwe_product(
209        module.n() as f64,
210        basek * digits,
211        0.5,
212        0.5,
213        0f64,
214        sigma * sigma,
215        0f64,
216        rank as f64,
217        k_ct,
218        k_ksk,
219    );
220
221    ct_glwe.assert_noise(module, &sk_out_prepared, &pt_want, max_noise + 0.5);
222}