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