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