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