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