1use backend::{FFT64, Module, ScalarZnxOps, ScratchOwned, Stats, VecZnxOps};
  2use sampling::source::Source;
  3
  4use crate::{
  5    AutomorphismKey, GLWECiphertextFourier, GLWEPlaintext, GLWESecret, GetRow, Infos, div_ceil,
  6    test_fft64::log2_std_noise_gglwe_product,
  7};
  8
  9#[test]
 10fn automorphism() {
 11    let log_n: usize = 8;
 12    let basek: usize = 12;
 13    let k_in: usize = 60;
 14    let k_out: usize = 60;
 15    let digits: usize = div_ceil(k_in, basek);
 16    let sigma: f64 = 3.2;
 17    (1..4).for_each(|rank| {
 18        (2..digits + 1).for_each(|di| {
 19            println!("test automorphism digits: {} rank: {}", di, rank);
 20            let k_apply: usize = (digits + di) * basek;
 21            test_automorphism(-1, 5, log_n, basek, di, k_in, k_out, k_apply, sigma, rank);
 22        });
 23    });
 24}
 25
 26#[test]
 27fn automorphism_inplace() {
 28    let log_n: usize = 8;
 29    let basek: usize = 12;
 30    let k_in: usize = 60;
 31    let digits: usize = div_ceil(k_in, basek);
 32    let sigma: f64 = 3.2;
 33    (1..4).for_each(|rank| {
 34        (2..digits + 1).for_each(|di| {
 35            println!("test automorphism digits: {} rank: {}", di, rank);
 36            let k_apply: usize = (digits + di) * basek;
 37            test_automorphism_inplace(-1, 5, log_n, basek, di, k_in, k_apply, sigma, rank);
 38        });
 39    });
 40}
 41
 42fn test_automorphism(
 43    p0: i64,
 44    p1: i64,
 45    log_n: usize,
 46    basek: usize,
 47    digits: usize,
 48    k_in: usize,
 49    k_out: usize,
 50    k_apply: usize,
 51    sigma: f64,
 52    rank: usize,
 53) {
 54    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
 55
 56    let digits_in: usize = 1;
 57
 58    let rows_in: usize = k_in / (basek * digits);
 59    let rows_apply: usize = div_ceil(k_in, basek * digits);
 60
 61    let mut auto_key_in: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::alloc(&module, basek, k_in, rows_in, digits_in, rank);
 62    let mut auto_key_out: AutomorphismKey<Vec<u8>, FFT64> =
 63        AutomorphismKey::alloc(&module, basek, k_out, rows_in, digits_in, rank);
 64    let mut auto_key_apply: AutomorphismKey<Vec<u8>, FFT64> =
 65        AutomorphismKey::alloc(&module, basek, k_apply, rows_apply, digits, rank);
 66
 67    let mut source_xs: Source = Source::new([0u8; 32]);
 68    let mut source_xe: Source = Source::new([0u8; 32]);
 69    let mut source_xa: Source = Source::new([0u8; 32]);
 70
 71    let mut scratch: ScratchOwned = ScratchOwned::new(
 72        AutomorphismKey::generate_from_sk_scratch_space(&module, basek, k_apply, rank)
 73            | GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k_out)
 74            | AutomorphismKey::automorphism_scratch_space(&module, basek, k_out, k_in, k_apply, digits, rank),
 75    );
 76
 77    let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
 78    sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
 79
 80    // gglwe_{s1}(s0) = s0 -> s1
 81    auto_key_in.generate_from_sk(
 82        &module,
 83        p0,
 84        &sk,
 85        &mut source_xa,
 86        &mut source_xe,
 87        sigma,
 88        scratch.borrow(),
 89    );
 90
 91    // gglwe_{s2}(s1) -> s1 -> s2
 92    auto_key_apply.generate_from_sk(
 93        &module,
 94        p1,
 95        &sk,
 96        &mut source_xa,
 97        &mut source_xe,
 98        sigma,
 99        scratch.borrow(),
100    );
101
102    // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
103    auto_key_out.automorphism(&module, &auto_key_in, &auto_key_apply, scratch.borrow());
104
105    let mut ct_glwe_dft: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_out, rank);
106    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_out);
107
108    let mut sk_auto: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
109    sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
110    (0..rank).for_each(|i| {
111        module.scalar_znx_automorphism(
112            module.galois_element_inv(p0 * p1),
113            &mut sk_auto.data,
114            i,
115            &sk.data,
116            i,
117        );
118    });
119
120    sk_auto.prep_fourier(&module);
121
122    (0..auto_key_out.rank_in()).for_each(|col_i| {
123        (0..auto_key_out.rows()).for_each(|row_i| {
124            auto_key_out.get_row(&module, row_i, col_i, &mut ct_glwe_dft);
125            ct_glwe_dft.decrypt(&module, &mut pt, &sk_auto, scratch.borrow());
126
127            module.vec_znx_sub_scalar_inplace(
128                &mut pt.data,
129                0,
130                (digits_in - 1) + row_i * digits_in,
131                &sk.data,
132                col_i,
133            );
134
135            let noise_have: f64 = pt.data.std(0, basek).log2();
136            let noise_want: f64 = log2_std_noise_gglwe_product(
137                module.n() as f64,
138                basek * digits,
139                0.5,
140                0.5,
141                0f64,
142                sigma * sigma,
143                0f64,
144                rank as f64,
145                k_in,
146                k_apply,
147            );
148
149            assert!(
150                (noise_have - noise_want).abs() <= 0.5,
151                "{} {}",
152                noise_have,
153                noise_want
154            );
155        });
156    });
157}
158
159fn test_automorphism_inplace(
160    p0: i64,
161    p1: i64,
162    log_n: usize,
163    basek: usize,
164    digits: usize,
165    k_in: usize,
166    k_apply: usize,
167    sigma: f64,
168    rank: usize,
169) {
170    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
171
172    let digits_in: usize = 1;
173
174    let rows_in: usize = k_in / (basek * digits);
175    let rows_apply: usize = div_ceil(k_in, basek * digits);
176
177    let mut auto_key: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::alloc(&module, basek, k_in, rows_in, digits_in, rank);
178    let mut auto_key_apply: AutomorphismKey<Vec<u8>, FFT64> =
179        AutomorphismKey::alloc(&module, basek, k_apply, rows_apply, digits, rank);
180
181    let mut source_xs: Source = Source::new([0u8; 32]);
182    let mut source_xe: Source = Source::new([0u8; 32]);
183    let mut source_xa: Source = Source::new([0u8; 32]);
184
185    let mut scratch: ScratchOwned = ScratchOwned::new(
186        AutomorphismKey::generate_from_sk_scratch_space(&module, basek, k_apply, rank)
187            | GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k_in)
188            | AutomorphismKey::automorphism_inplace_scratch_space(&module, basek, k_in, k_apply, digits, rank),
189    );
190
191    let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
192    sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
193
194    // gglwe_{s1}(s0) = s0 -> s1
195    auto_key.generate_from_sk(
196        &module,
197        p0,
198        &sk,
199        &mut source_xa,
200        &mut source_xe,
201        sigma,
202        scratch.borrow(),
203    );
204
205    // gglwe_{s2}(s1) -> s1 -> s2
206    auto_key_apply.generate_from_sk(
207        &module,
208        p1,
209        &sk,
210        &mut source_xa,
211        &mut source_xe,
212        sigma,
213        scratch.borrow(),
214    );
215
216    // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
217    auto_key.automorphism_inplace(&module, &auto_key_apply, scratch.borrow());
218
219    let mut ct_glwe_dft: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_in, rank);
220    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_in);
221
222    let mut sk_auto: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
223    sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
224    (0..rank).for_each(|i| {
225        module.scalar_znx_automorphism(
226            module.galois_element_inv(p0 * p1),
227            &mut sk_auto.data,
228            i,
229            &sk.data,
230            i,
231        );
232    });
233
234    sk_auto.prep_fourier(&module);
235
236    (0..auto_key.rank_in()).for_each(|col_i| {
237        (0..auto_key.rows()).for_each(|row_i| {
238            auto_key.get_row(&module, row_i, col_i, &mut ct_glwe_dft);
239
240            ct_glwe_dft.decrypt(&module, &mut pt, &sk_auto, scratch.borrow());
241            module.vec_znx_sub_scalar_inplace(
242                &mut pt.data,
243                0,
244                (digits_in - 1) + row_i * digits_in,
245                &sk.data,
246                col_i,
247            );
248
249            let noise_have: f64 = pt.data.std(0, basek).log2();
250            let noise_want: f64 = log2_std_noise_gglwe_product(
251                module.n() as f64,
252                basek * digits,
253                0.5,
254                0.5,
255                0f64,
256                sigma * sigma,
257                0f64,
258                rank as f64,
259                k_in,
260                k_apply,
261            );
262
263            assert!(
264                (noise_have - noise_want).abs() <= 0.5,
265                "{} {}",
266                noise_have,
267                noise_want
268            );
269        });
270    });
271}