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