1use backend::{FFT64, FillUniform, Module, ScratchOwned, Stats, VecZnxOps};
  2use sampling::source::Source;
  3
  4use crate::{
  5    FourierGLWESecret, GLWECiphertext, GLWEPlaintext, GLWESecret, GLWESwitchingKey, Infos, noise::log2_std_noise_gglwe_product,
  6};
  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                let k_out: usize = k_ksk; // better capture noise
 19                println!(
 20                    "test keyswitch digits: {} rank_in: {} rank_out: {}",
 21                    di, rank_in, rank_out
 22                );
 23                test_keyswitch(log_n, basek, k_out, k_in, 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_keyswitch_inplace(log_n, basek, k_ct, k_ksk, di, rank, 3.2);
 40        });
 41    });
 42}
 43
 44fn test_keyswitch(
 45    log_n: usize,
 46    basek: usize,
 47    k_out: usize,
 48    k_in: 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_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_in, rank_in);
 62    let mut ct_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_out, rank_out);
 63    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_in);
 64    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_out);
 65
 66    let mut source_xs: Source = Source::new([0u8; 32]);
 67    let mut source_xe: Source = Source::new([0u8; 32]);
 68    let mut source_xa: Source = Source::new([0u8; 32]);
 69
 70    pt_want
 71        .data
 72        .fill_uniform(basek, 0, pt_want.size(), &mut source_xa);
 73
 74    let mut scratch: ScratchOwned = ScratchOwned::new(
 75        GLWESwitchingKey::encrypt_sk_scratch_space(&module, basek, ksk.k(), rank_out)
 76            | GLWECiphertext::decrypt_scratch_space(&module, basek, ct_out.k())
 77            | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct_in.k())
 78            | GLWECiphertext::keyswitch_scratch_space(
 79                &module,
 80                basek,
 81                ct_out.k(),
 82                ct_in.k(),
 83                ksk.k(),
 84                digits,
 85                rank_in,
 86                rank_out,
 87            ),
 88    );
 89
 90    let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank_in);
 91    sk_in.fill_ternary_prob(0.5, &mut source_xs);
 92    let sk_in_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::from(&module, &sk_in);
 93
 94    let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank_out);
 95    sk_out.fill_ternary_prob(0.5, &mut source_xs);
 96    let sk_out_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::from(&module, &sk_out);
 97
 98    ksk.generate_from_sk(
 99        &module,
100        &sk_in,
101        &sk_out_dft,
102        &mut source_xa,
103        &mut source_xe,
104        sigma,
105        scratch.borrow(),
106    );
107
108    ct_in.encrypt_sk(
109        &module,
110        &pt_want,
111        &sk_in_dft,
112        &mut source_xa,
113        &mut source_xe,
114        sigma,
115        scratch.borrow(),
116    );
117
118    ct_out.keyswitch(&module, &ct_in, &ksk, scratch.borrow());
119    ct_out.decrypt(&module, &mut pt_have, &sk_out_dft, scratch.borrow());
120
121    module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
122
123    let noise_have: f64 = pt_have.data.std(0, basek).log2();
124    let noise_want: f64 = log2_std_noise_gglwe_product(
125        module.n() as f64,
126        basek * digits,
127        0.5,
128        0.5,
129        0f64,
130        sigma * sigma,
131        0f64,
132        rank_in as f64,
133        k_in,
134        k_ksk,
135    );
136
137    println!("{} vs. {}", noise_have, noise_want);
138
139    assert!(
140        (noise_have - noise_want).abs() <= 0.5,
141        "{} {}",
142        noise_have,
143        noise_want
144    );
145}
146
147fn test_keyswitch_inplace(log_n: usize, basek: usize, k_ct: usize, k_ksk: usize, digits: usize, rank: usize, sigma: f64) {
148    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
149
150    let rows: usize = k_ct.div_ceil(basek * digits);
151
152    let mut ct_grlwe: GLWESwitchingKey<Vec<u8>, FFT64> = GLWESwitchingKey::alloc(&module, basek, k_ksk, rows, digits, rank, rank);
153    let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct, rank);
154    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
155    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
156
157    let mut source_xs: Source = Source::new([0u8; 32]);
158    let mut source_xe: Source = Source::new([0u8; 32]);
159    let mut source_xa: Source = Source::new([0u8; 32]);
160
161    pt_want
162        .data
163        .fill_uniform(basek, 0, pt_want.size(), &mut source_xa);
164
165    let mut scratch: ScratchOwned = ScratchOwned::new(
166        GLWESwitchingKey::encrypt_sk_scratch_space(&module, basek, ct_grlwe.k(), rank)
167            | GLWECiphertext::decrypt_scratch_space(&module, basek, ct_glwe.k())
168            | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct_glwe.k())
169            | GLWECiphertext::keyswitch_inplace_scratch_space(&module, basek, ct_glwe.k(), ct_grlwe.k(), digits, rank),
170    );
171
172    let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
173    sk_in.fill_ternary_prob(0.5, &mut source_xs);
174    let sk_in_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::from(&module, &sk_in);
175
176    let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
177    sk_out.fill_ternary_prob(0.5, &mut source_xs);
178    let sk_out_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::from(&module, &sk_out);
179
180    ct_grlwe.generate_from_sk(
181        &module,
182        &sk_in,
183        &sk_out_dft,
184        &mut source_xa,
185        &mut source_xe,
186        sigma,
187        scratch.borrow(),
188    );
189
190    ct_glwe.encrypt_sk(
191        &module,
192        &pt_want,
193        &sk_in_dft,
194        &mut source_xa,
195        &mut source_xe,
196        sigma,
197        scratch.borrow(),
198    );
199
200    ct_glwe.keyswitch_inplace(&module, &ct_grlwe, scratch.borrow());
201
202    ct_glwe.decrypt(&module, &mut pt_have, &sk_out_dft, scratch.borrow());
203
204    module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
205
206    let noise_have: f64 = pt_have.data.std(0, basek).log2();
207    let noise_want: f64 = log2_std_noise_gglwe_product(
208        module.n() as f64,
209        basek * digits,
210        0.5,
211        0.5,
212        0f64,
213        sigma * sigma,
214        0f64,
215        rank as f64,
216        k_ct,
217        k_ksk,
218    );
219
220    assert!(
221        (noise_have - noise_want).abs() <= 0.5,
222        "{} {}",
223        noise_have,
224        noise_want
225    );
226}