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