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