1use backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApply, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
  4        VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace,
  5        VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc, VecZnxBigAllocBytes, VecZnxBigNormalize,
  6        VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftAddInplace, VecZnxDftAlloc, VecZnxDftAllocBytes, VecZnxDftCopy,
  7        VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxDftToVecZnxBigTmpA, VecZnxFillUniform, VecZnxNormalize,
  8        VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, VecZnxSwithcDegree, VmpApply,
  9        VmpApplyAdd, VmpApplyTmpBytes, VmpPMatAlloc, VmpPrepare,
 10    },
 11    layouts::{Backend, Module, ScalarZnx, ScratchOwned},
 12    oep::{
 13        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
 14        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
 15    },
 16};
 17use sampling::source::Source;
 18
 19use crate::{
 20    layouts::{
 21        GGLWEAutomorphismKey, GGLWETensorKey, GGSWCiphertext, GLWESecret,
 22        prepared::{GGLWEAutomorphismKeyPrepared, GGLWETensorKeyPrepared, GLWESecretPrepared, Prepare, PrepareAlloc},
 23    },
 24    noise::noise_ggsw_keyswitch,
 25};
 26
 27#[allow(clippy::too_many_arguments)]
 28pub fn test_ggsw_automorphism<B>(
 29    p: i64,
 30    module: &Module<B>,
 31    basek: usize,
 32    k_out: usize,
 33    k_in: usize,
 34    k_ksk: usize,
 35    k_tsk: usize,
 36    digits: usize,
 37    rank: usize,
 38    sigma: f64,
 39) where
 40    Module<B>: VecZnxDftAllocBytes
 41        + VecZnxBigAllocBytes
 42        + VecZnxDftFromVecZnx<B>
 43        + SvpApplyInplace<B>
 44        + VecZnxDftToVecZnxBigConsume<B>
 45        + VecZnxBigAddInplace<B>
 46        + VecZnxBigAddSmallInplace<B>
 47        + VecZnxBigNormalize<B>
 48        + VecZnxNormalizeTmpBytes
 49        + VecZnxBigAlloc<B>
 50        + VecZnxDftAlloc<B>
 51        + VecZnxBigNormalizeTmpBytes
 52        + VecZnxDftToVecZnxBigTmpA<B>
 53        + SvpPrepare<B>
 54        + SvpPPolAllocBytes
 55        + SvpPPolAlloc<B>
 56        + VecZnxAddScalarInplace
 57        + VecZnxCopy
 58        + VecZnxSubABInplace
 59        + VmpPMatAlloc<B>
 60        + VmpPrepare<B>
 61        + VmpApplyTmpBytes
 62        + VmpApply<B>
 63        + VmpApplyAdd<B>
 64        + VecZnxDftCopy<B>
 65        + VecZnxDftAddInplace<B>
 66        + VecZnxFillUniform
 67        + VecZnxAddInplace
 68        + VecZnxNormalizeInplace<B>
 69        + VecZnxAddNormal
 70        + VecZnxNormalize<B>
 71        + VecZnxSub
 72        + SvpApply<B>
 73        + VecZnxSwithcDegree
 74        + VecZnxAutomorphismInplace
 75        + VecZnxAutomorphism,
 76    B: Backend
 77        + TakeVecZnxDftImpl<B>
 78        + TakeVecZnxBigImpl<B>
 79        + ScratchOwnedAllocImpl<B>
 80        + ScratchOwnedBorrowImpl<B>
 81        + ScratchAvailableImpl<B>
 82        + TakeScalarZnxImpl<B>
 83        + TakeVecZnxImpl<B>
 84        + VecZnxDftAllocBytesImpl<B>
 85        + VecZnxBigAllocBytesImpl<B>
 86        + TakeSvpPPolImpl<B>,
 87{
 88    let n: usize = module.n();
 89    let rows: usize = k_in.div_ceil(basek * digits);
 90    let rows_in: usize = k_in.div_euclid(basek * digits);
 91
 92    let digits_in: usize = 1;
 93
 94    let mut ct_in: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_in, rows_in, digits_in, rank);
 95    let mut ct_out: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_out, rows_in, digits_in, rank);
 96    let mut tensor_key: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_tsk, rows, digits, rank);
 97    let mut auto_key: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
 98    let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
 99
100    let mut source_xs: Source = Source::new([0u8; 32]);
101    let mut source_xe: Source = Source::new([0u8; 32]);
102    let mut source_xa: Source = Source::new([0u8; 32]);
103
104    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
105        GGSWCiphertext::encrypt_sk_scratch_space(module, n, basek, k_in, rank)
106            | GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank)
107            | GGLWETensorKey::encrypt_sk_scratch_space(module, n, basek, k_tsk, rank)
108            | GGSWCiphertext::automorphism_scratch_space(
109                module, n, basek, k_out, k_in, k_ksk, digits, k_tsk, digits, rank,
110            ),
111    );
112
113    let var_xs: f64 = 0.5;
114
115    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
116    sk.fill_ternary_prob(var_xs, &mut source_xs);
117    let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
118
119    auto_key.encrypt_sk(
120        module,
121        p,
122        &sk,
123        &mut source_xa,
124        &mut source_xe,
125        sigma,
126        scratch.borrow(),
127    );
128    tensor_key.encrypt_sk(
129        module,
130        &sk,
131        &mut source_xa,
132        &mut source_xe,
133        sigma,
134        scratch.borrow(),
135    );
136
137    pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
138
139    ct_in.encrypt_sk(
140        module,
141        &pt_scalar,
142        &sk_prepared,
143        &mut source_xa,
144        &mut source_xe,
145        sigma,
146        scratch.borrow(),
147    );
148
149    let mut auto_key_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, B> =
150        GGLWEAutomorphismKeyPrepared::alloc(module, n, basek, k_ksk, rows, digits, rank);
151    auto_key_prepared.prepare(module, &auto_key, scratch.borrow());
152
153    let mut tsk_prepared: GGLWETensorKeyPrepared<Vec<u8>, B> =
154        GGLWETensorKeyPrepared::alloc(module, n, basek, k_tsk, rows, digits, rank);
155    tsk_prepared.prepare(module, &tensor_key, scratch.borrow());
156
157    ct_out.automorphism(
158        module,
159        &ct_in,
160        &auto_key_prepared,
161        &tsk_prepared,
162        scratch.borrow(),
163    );
164
165    module.vec_znx_automorphism_inplace(p, &mut pt_scalar.as_vec_znx_mut(), 0);
166
167    let max_noise = |col_j: usize| -> f64 {
168        noise_ggsw_keyswitch(
169            n as f64,
170            basek * digits,
171            col_j,
172            var_xs,
173            0f64,
174            sigma * sigma,
175            0f64,
176            rank as f64,
177            k_in,
178            k_ksk,
179            k_tsk,
180        ) + 0.5
181    };
182
183    ct_out.assert_noise(module, &sk_prepared, &pt_scalar, max_noise);
184}
185
186#[allow(clippy::too_many_arguments)]
187pub fn test_ggsw_automorphism_inplace<B>(
188    p: i64,
189    module: &Module<B>,
190    basek: usize,
191    k_ct: usize,
192    k_ksk: usize,
193    k_tsk: usize,
194    digits: usize,
195    rank: usize,
196    sigma: f64,
197) where
198    Module<B>: VecZnxDftAllocBytes
199        + VecZnxBigAllocBytes
200        + VecZnxDftFromVecZnx<B>
201        + SvpApplyInplace<B>
202        + VecZnxDftToVecZnxBigConsume<B>
203        + VecZnxBigAddInplace<B>
204        + VecZnxBigAddSmallInplace<B>
205        + VecZnxBigNormalize<B>
206        + VecZnxNormalizeTmpBytes
207        + VecZnxBigAlloc<B>
208        + VecZnxDftAlloc<B>
209        + VecZnxBigNormalizeTmpBytes
210        + VecZnxDftToVecZnxBigTmpA<B>
211        + SvpPrepare<B>
212        + SvpPPolAllocBytes
213        + SvpPPolAlloc<B>
214        + VecZnxAddScalarInplace
215        + VecZnxCopy
216        + VecZnxSubABInplace
217        + VmpPMatAlloc<B>
218        + VmpPrepare<B>
219        + VmpApplyTmpBytes
220        + VmpApply<B>
221        + VmpApplyAdd<B>
222        + VecZnxBigAddSmallInplace<B>
223        + VecZnxDftCopy<B>
224        + VecZnxDftAddInplace<B>
225        + VecZnxAddInplace
226        + VecZnxNormalizeInplace<B>
227        + VecZnxAddNormal
228        + VecZnxNormalize<B>
229        + VecZnxSub
230        + VecZnxFillUniform
231        + SvpApply<B>
232        + VecZnxSwithcDegree
233        + VecZnxAutomorphismInplace
234        + VecZnxAutomorphism,
235    B: Backend
236        + TakeVecZnxDftImpl<B>
237        + TakeVecZnxBigImpl<B>
238        + ScratchOwnedAllocImpl<B>
239        + ScratchOwnedBorrowImpl<B>
240        + ScratchAvailableImpl<B>
241        + TakeScalarZnxImpl<B>
242        + TakeVecZnxImpl<B>
243        + VecZnxDftAllocBytesImpl<B>
244        + VecZnxBigAllocBytesImpl<B>
245        + TakeSvpPPolImpl<B>,
246{
247    let n: usize = module.n();
248    let rows: usize = k_ct.div_ceil(digits * basek);
249    let rows_in: usize = k_ct.div_euclid(basek * digits);
250    let digits_in: usize = 1;
251
252    let mut ct: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_ct, rows_in, digits_in, rank);
253    let mut tensor_key: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_tsk, rows, digits, rank);
254    let mut auto_key: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
255    let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
256
257    let mut source_xs: Source = Source::new([0u8; 32]);
258    let mut source_xe: Source = Source::new([0u8; 32]);
259    let mut source_xa: Source = Source::new([0u8; 32]);
260
261    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
262        GGSWCiphertext::encrypt_sk_scratch_space(module, n, basek, k_ct, rank)
263            | GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank)
264            | GGLWETensorKey::encrypt_sk_scratch_space(module, n, basek, k_tsk, rank)
265            | GGSWCiphertext::automorphism_inplace_scratch_space(module, n, basek, k_ct, k_ksk, digits, k_tsk, digits, rank),
266    );
267
268    let var_xs: f64 = 0.5;
269
270    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
271    sk.fill_ternary_prob(var_xs, &mut source_xs);
272    let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
273
274    auto_key.encrypt_sk(
275        module,
276        p,
277        &sk,
278        &mut source_xa,
279        &mut source_xe,
280        sigma,
281        scratch.borrow(),
282    );
283    tensor_key.encrypt_sk(
284        module,
285        &sk,
286        &mut source_xa,
287        &mut source_xe,
288        sigma,
289        scratch.borrow(),
290    );
291
292    pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
293
294    ct.encrypt_sk(
295        module,
296        &pt_scalar,
297        &sk_prepared,
298        &mut source_xa,
299        &mut source_xe,
300        sigma,
301        scratch.borrow(),
302    );
303
304    let mut auto_key_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, B> =
305        GGLWEAutomorphismKeyPrepared::alloc(module, n, basek, k_ksk, rows, digits, rank);
306    auto_key_prepared.prepare(module, &auto_key, scratch.borrow());
307
308    let mut tsk_prepared: GGLWETensorKeyPrepared<Vec<u8>, B> =
309        GGLWETensorKeyPrepared::alloc(module, n, basek, k_tsk, rows, digits, rank);
310    tsk_prepared.prepare(module, &tensor_key, scratch.borrow());
311
312    ct.automorphism_inplace(module, &auto_key_prepared, &tsk_prepared, scratch.borrow());
313
314    module.vec_znx_automorphism_inplace(p, &mut pt_scalar.as_vec_znx_mut(), 0);
315
316    let max_noise = |col_j: usize| -> f64 {
317        noise_ggsw_keyswitch(
318            n as f64,
319            basek * digits,
320            col_j,
321            var_xs,
322            0f64,
323            sigma * sigma,
324            0f64,
325            rank as f64,
326            k_ct,
327            k_ksk,
328            k_tsk,
329        ) + 0.5
330    };
331
332    ct.assert_noise(module, &sk_prepared, &pt_scalar, max_noise);
333}