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