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