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