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