1use backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxCopy,
  4        VecZnxStd, VecZnxSubScalarInplace, VecZnxSwithcDegree, 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        GGLWEAutomorphismKey, GLWEPlaintext, GLWESecret, Infos,
 17        prepared::{GGLWEAutomorphismKeyExec, GLWESecretExec},
 18    },
 19    noise::log2_std_noise_gglwe_product,
 20    trait_families::{GLWEDecryptFamily, GLWEKeyswitchFamily},
 21};
 22
 23use crate::trait_families::{GGLWEAutomorphismKeyEncryptSkFamily, GLWESecretExecModuleFamily};
 24
 25pub fn test_gglwe_automorphism_key_automorphism<B: Backend>(
 26    module: &Module<B>,
 27    p0: i64,
 28    p1: i64,
 29    basek: usize,
 30    digits: usize,
 31    k_in: usize,
 32    k_out: usize,
 33    k_apply: usize,
 34    sigma: f64,
 35    rank: usize,
 36) where
 37    Module<B>: GGLWEAutomorphismKeyEncryptSkFamily<B>
 38        + GLWESecretExecModuleFamily<B>
 39        + GLWEKeyswitchFamily<B>
 40        + VecZnxAutomorphism
 41        + VecZnxSwithcDegree
 42        + VecZnxAddScalarInplace
 43        + VecZnxAutomorphism
 44        + VecZnxAutomorphismInplace
 45        + GLWEDecryptFamily<B>
 46        + VecZnxSubScalarInplace
 47        + VecZnxStd
 48        + VecZnxCopy
 49        + VmpPMatAlloc<B>
 50        + VmpPMatPrepare<B>,
 51    B: ScratchOwnedAllocImpl<B>
 52        + ScratchOwnedBorrowImpl<B>
 53        + ScratchAvailableImpl<B>
 54        + TakeScalarZnxImpl<B>
 55        + TakeVecZnxDftImpl<B>
 56        + TakeVecZnxImpl<B>
 57        + TakeSvpPPolImpl<B>
 58        + TakeVecZnxBigImpl<B>,
 59{
 60    let n: usize = module.n();
 61    let digits_in: usize = 1;
 62
 63    let rows_in: usize = k_in / (basek * digits);
 64    let rows_apply: usize = k_in.div_ceil(basek * digits);
 65
 66    let mut auto_key_in: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_in, rows_in, digits_in, rank);
 67    let mut auto_key_out: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_out, rows_in, digits_in, rank);
 68    let mut auto_key_apply: GGLWEAutomorphismKey<Vec<u8>> =
 69        GGLWEAutomorphismKey::alloc(n, basek, k_apply, rows_apply, digits, rank);
 70
 71    let mut source_xs: Source = Source::new([0u8; 32]);
 72    let mut source_xe: Source = Source::new([0u8; 32]);
 73    let mut source_xa: Source = Source::new([0u8; 32]);
 74
 75    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 76        GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_apply, rank)
 77            | GGLWEAutomorphismKey::automorphism_scratch_space(module, n, basek, k_out, k_in, k_apply, digits, rank),
 78    );
 79
 80    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
 81    sk.fill_ternary_prob(0.5, &mut source_xs);
 82
 83    // gglwe_{s1}(s0) = s0 -> s1
 84    auto_key_in.encrypt_sk(
 85        module,
 86        p0,
 87        &sk,
 88        &mut source_xa,
 89        &mut source_xe,
 90        sigma,
 91        scratch.borrow(),
 92    );
 93
 94    // gglwe_{s2}(s1) -> s1 -> s2
 95    auto_key_apply.encrypt_sk(
 96        module,
 97        p1,
 98        &sk,
 99        &mut source_xa,
100        &mut source_xe,
101        sigma,
102        scratch.borrow(),
103    );
104
105    let mut auto_key_apply_exec: GGLWEAutomorphismKeyExec<Vec<u8>, B> =
106        GGLWEAutomorphismKeyExec::alloc(module, n, basek, k_apply, rows_apply, digits, rank);
107
108    auto_key_apply_exec.prepare(module, &auto_key_apply, scratch.borrow());
109
110    // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
111    auto_key_out.automorphism(module, &auto_key_in, &auto_key_apply_exec, scratch.borrow());
112
113    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_out);
114
115    let mut sk_auto: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
116    sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
117    (0..rank).for_each(|i| {
118        module.vec_znx_automorphism(
119            module.galois_element_inv(p0 * p1),
120            &mut sk_auto.data.as_vec_znx_mut(),
121            i,
122            &sk.data.as_vec_znx(),
123            i,
124        );
125    });
126
127    let sk_auto_dft: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_auto);
128
129    (0..auto_key_out.rank_in()).for_each(|col_i| {
130        (0..auto_key_out.rows()).for_each(|row_i| {
131            auto_key_out
132                .at(row_i, col_i)
133                .decrypt(module, &mut pt, &sk_auto_dft, scratch.borrow());
134
135            module.vec_znx_sub_scalar_inplace(
136                &mut pt.data,
137                0,
138                (digits_in - 1) + row_i * digits_in,
139                &sk.data,
140                col_i,
141            );
142
143            let noise_have: f64 = module.vec_znx_std(basek, &pt.data, 0).log2();
144            let noise_want: f64 = log2_std_noise_gglwe_product(
145                n as f64,
146                basek * digits,
147                0.5,
148                0.5,
149                0f64,
150                sigma * sigma,
151                0f64,
152                rank as f64,
153                k_out,
154                k_apply,
155            );
156
157            assert!(
158                noise_have < noise_want + 0.5,
159                "{} {}",
160                noise_have,
161                noise_want
162            );
163        });
164    });
165}
166
167pub fn test_gglwe_automorphism_key_automorphism_inplace<B: Backend>(
168    module: &Module<B>,
169    p0: i64,
170    p1: i64,
171    basek: usize,
172    digits: usize,
173    k_in: usize,
174    k_apply: usize,
175    sigma: f64,
176    rank: usize,
177) where
178    Module<B>: GGLWEAutomorphismKeyEncryptSkFamily<B>
179        + GLWESecretExecModuleFamily<B>
180        + GLWEKeyswitchFamily<B>
181        + VecZnxAutomorphism
182        + VecZnxSwithcDegree
183        + VecZnxAddScalarInplace
184        + VecZnxAutomorphism
185        + VecZnxAutomorphismInplace
186        + GLWEDecryptFamily<B>
187        + VecZnxSubScalarInplace
188        + VecZnxStd
189        + VecZnxCopy
190        + VmpPMatAlloc<B>
191        + VmpPMatPrepare<B>,
192    B: ScratchOwnedAllocImpl<B>
193        + ScratchOwnedBorrowImpl<B>
194        + ScratchAvailableImpl<B>
195        + TakeScalarZnxImpl<B>
196        + TakeVecZnxDftImpl<B>
197        + TakeVecZnxImpl<B>
198        + TakeSvpPPolImpl<B>
199        + TakeVecZnxBigImpl<B>,
200{
201    let n: usize = module.n();
202    let digits_in: usize = 1;
203
204    let rows_in: usize = k_in / (basek * digits);
205    let rows_apply: usize = k_in.div_ceil(basek * digits);
206
207    let mut auto_key: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_in, rows_in, digits_in, rank);
208    let mut auto_key_apply: GGLWEAutomorphismKey<Vec<u8>> =
209        GGLWEAutomorphismKey::alloc(n, basek, k_apply, rows_apply, digits, rank);
210
211    let mut source_xs: Source = Source::new([0u8; 32]);
212    let mut source_xe: Source = Source::new([0u8; 32]);
213    let mut source_xa: Source = Source::new([0u8; 32]);
214
215    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
216        GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_apply, rank)
217            | GGLWEAutomorphismKey::automorphism_inplace_scratch_space(module, n, basek, k_in, k_apply, digits, rank),
218    );
219
220    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
221    sk.fill_ternary_prob(0.5, &mut source_xs);
222
223    // gglwe_{s1}(s0) = s0 -> s1
224    auto_key.encrypt_sk(
225        module,
226        p0,
227        &sk,
228        &mut source_xa,
229        &mut source_xe,
230        sigma,
231        scratch.borrow(),
232    );
233
234    // gglwe_{s2}(s1) -> s1 -> s2
235    auto_key_apply.encrypt_sk(
236        module,
237        p1,
238        &sk,
239        &mut source_xa,
240        &mut source_xe,
241        sigma,
242        scratch.borrow(),
243    );
244
245    let mut auto_key_apply_exec: GGLWEAutomorphismKeyExec<Vec<u8>, B> =
246        GGLWEAutomorphismKeyExec::alloc(module, n, basek, k_apply, rows_apply, digits, rank);
247
248    auto_key_apply_exec.prepare(module, &auto_key_apply, scratch.borrow());
249
250    // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
251    auto_key.automorphism_inplace(module, &auto_key_apply_exec, scratch.borrow());
252
253    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_in);
254
255    let mut sk_auto: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
256    sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
257
258    (0..rank).for_each(|i| {
259        module.vec_znx_automorphism(
260            module.galois_element_inv(p0 * p1),
261            &mut sk_auto.data.as_vec_znx_mut(),
262            i,
263            &sk.data.as_vec_znx(),
264            i,
265        );
266    });
267
268    let sk_auto_dft: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_auto);
269
270    (0..auto_key.rank_in()).for_each(|col_i| {
271        (0..auto_key.rows()).for_each(|row_i| {
272            auto_key
273                .at(row_i, col_i)
274                .decrypt(module, &mut pt, &sk_auto_dft, scratch.borrow());
275            module.vec_znx_sub_scalar_inplace(
276                &mut pt.data,
277                0,
278                (digits_in - 1) + row_i * digits_in,
279                &sk.data,
280                col_i,
281            );
282
283            let noise_have: f64 = module.vec_znx_std(basek, &pt.data, 0).log2();
284            let noise_want: f64 = log2_std_noise_gglwe_product(
285                n as f64,
286                basek * digits,
287                0.5,
288                0.5,
289                0f64,
290                sigma * sigma,
291                0f64,
292                rank as f64,
293                k_in,
294                k_apply,
295            );
296
297            assert!(
298                noise_have < noise_want + 0.5,
299                "{} {}",
300                noise_have,
301                noise_want
302            );
303        });
304    });
305}