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