1use backend::{FFT64, FillUniform, Module, ScratchOwned, Stats, VecZnxOps};
  2
  3use sampling::source::Source;
  4
  5use crate::{
  6    FourierGLWESecret, GLWEAutomorphismKey, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos, noise::log2_std_noise_gglwe_product,
  7};
  8
  9#[test]
 10fn apply_inplace() {
 11    let log_n: usize = 8;
 12    let basek: usize = 12;
 13    let k_ct: usize = 60;
 14    let digits: usize = k_ct.div_ceil(basek);
 15    (1..4).for_each(|rank| {
 16        (1..digits + 1).for_each(|di| {
 17            let k_ksk: usize = k_ct + basek * di;
 18            println!("test automorphism_inplace digits: {} rank: {}", di, rank);
 19            test_automorphism_inplace(log_n, basek, -5, k_ct, k_ksk, di, rank, 3.2);
 20        });
 21    });
 22}
 23
 24#[test]
 25fn apply() {
 26    let log_n: usize = 8;
 27    let basek: usize = 12;
 28    let k_in: usize = 60;
 29    let digits: usize = k_in.div_ceil(basek);
 30    (1..4).for_each(|rank| {
 31        (1..digits + 1).for_each(|di| {
 32            let k_ksk: usize = k_in + basek * di;
 33            let k_out: usize = k_ksk; // Better capture noise.
 34            println!("test automorphism digits: {} rank: {}", di, rank);
 35            test_automorphism(log_n, basek, -5, k_out, k_in, k_ksk, di, rank, 3.2);
 36        })
 37    });
 38}
 39
 40fn test_automorphism(
 41    log_n: usize,
 42    basek: usize,
 43    p: i64,
 44    k_out: usize,
 45    k_in: usize,
 46    k_ksk: usize,
 47    digits: usize,
 48    rank: usize,
 49    sigma: f64,
 50) {
 51    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
 52
 53    let rows: usize = k_in.div_ceil(basek * digits);
 54
 55    let mut autokey: GLWEAutomorphismKey<Vec<u8>, FFT64> = GLWEAutomorphismKey::alloc(&module, basek, k_ksk, rows, digits, rank);
 56    let mut ct_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_in, rank);
 57    let mut ct_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_out, rank);
 58    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_in);
 59    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_out);
 60
 61    let mut source_xs: Source = Source::new([0u8; 32]);
 62    let mut source_xe: Source = Source::new([0u8; 32]);
 63    let mut source_xa: Source = Source::new([0u8; 32]);
 64
 65    pt_want
 66        .data
 67        .fill_uniform(basek, 0, pt_want.size(), &mut source_xa);
 68
 69    let mut scratch: ScratchOwned = ScratchOwned::new(
 70        GLWEAutomorphismKey::generate_from_sk_scratch_space(&module, basek, autokey.k(), rank)
 71            | GLWECiphertext::decrypt_scratch_space(&module, basek, ct_out.k())
 72            | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct_in.k())
 73            | GLWECiphertext::automorphism_scratch_space(
 74                &module,
 75                basek,
 76                ct_out.k(),
 77                ct_in.k(),
 78                autokey.k(),
 79                digits,
 80                rank,
 81            ),
 82    );
 83
 84    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
 85    sk.fill_ternary_prob(0.5, &mut source_xs);
 86    let sk_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::from(&module, &sk);
 87
 88    autokey.generate_from_sk(
 89        &module,
 90        p,
 91        &sk,
 92        &mut source_xa,
 93        &mut source_xe,
 94        sigma,
 95        scratch.borrow(),
 96    );
 97
 98    ct_in.encrypt_sk(
 99        &module,
100        &pt_want,
101        &sk_dft,
102        &mut source_xa,
103        &mut source_xe,
104        sigma,
105        scratch.borrow(),
106    );
107
108    ct_out.automorphism(&module, &ct_in, &autokey, scratch.borrow());
109    ct_out.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
110    module.vec_znx_automorphism_inplace(p, &mut pt_want.data, 0);
111    module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
112    module.vec_znx_normalize_inplace(basek, &mut pt_have.data, 0, scratch.borrow());
113
114    let noise_have: f64 = pt_have.data.std(0, basek).log2();
115
116    println!("{}", noise_have);
117
118    let noise_want: f64 = log2_std_noise_gglwe_product(
119        module.n() as f64,
120        basek * digits,
121        0.5,
122        0.5,
123        0f64,
124        sigma * sigma,
125        0f64,
126        rank as f64,
127        k_in,
128        k_ksk,
129    );
130
131    assert!(
132        noise_have <= noise_want + 1.0,
133        "{} {}",
134        noise_have,
135        noise_want
136    );
137}
138
139fn test_automorphism_inplace(
140    log_n: usize,
141    basek: usize,
142    p: i64,
143    k_ct: usize,
144    k_ksk: usize,
145    digits: usize,
146    rank: usize,
147    sigma: f64,
148) {
149    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
150
151    let rows: usize = k_ct.div_ceil(basek * digits);
152
153    let mut autokey: GLWEAutomorphismKey<Vec<u8>, FFT64> = GLWEAutomorphismKey::alloc(&module, basek, k_ksk, rows, digits, rank);
154    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct, rank);
155    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
156    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
157
158    let mut source_xs: Source = Source::new([0u8; 32]);
159    let mut source_xe: Source = Source::new([0u8; 32]);
160    let mut source_xa: Source = Source::new([0u8; 32]);
161
162    pt_want
163        .data
164        .fill_uniform(basek, 0, pt_want.size(), &mut source_xa);
165
166    let mut scratch: ScratchOwned = ScratchOwned::new(
167        GLWEAutomorphismKey::generate_from_sk_scratch_space(&module, basek, autokey.k(), rank)
168            | GLWECiphertext::decrypt_scratch_space(&module, basek, ct.k())
169            | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct.k())
170            | GLWECiphertext::automorphism_inplace_scratch_space(&module, basek, ct.k(), autokey.k(), digits, rank),
171    );
172
173    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
174    sk.fill_ternary_prob(0.5, &mut source_xs);
175    let sk_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::from(&module, &sk);
176
177    autokey.generate_from_sk(
178        &module,
179        p,
180        &sk,
181        &mut source_xa,
182        &mut source_xe,
183        sigma,
184        scratch.borrow(),
185    );
186
187    ct.encrypt_sk(
188        &module,
189        &pt_want,
190        &sk_dft,
191        &mut source_xa,
192        &mut source_xe,
193        sigma,
194        scratch.borrow(),
195    );
196
197    ct.automorphism_inplace(&module, &autokey, scratch.borrow());
198    ct.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
199    module.vec_znx_automorphism_inplace(p, &mut pt_want.data, 0);
200    module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
201    module.vec_znx_normalize_inplace(basek, &mut pt_have.data, 0, scratch.borrow());
202
203    let noise_have: f64 = pt_have.data.std(0, basek).log2();
204    let noise_want: f64 = log2_std_noise_gglwe_product(
205        module.n() as f64,
206        basek * digits,
207        0.5,
208        0.5,
209        0f64,
210        sigma * sigma,
211        0f64,
212        rank as f64,
213        k_ct,
214        k_ksk,
215    );
216
217    assert!(
218        (noise_have - noise_want).abs() <= 0.5,
219        "{} {}",
220        noise_have,
221        noise_want
222    );
223}