arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/poulpy-core/src/tests/test_suite/keyswitch/gglwe_ct.rs
poulpy-core/src/tests/test_suite/keyswitch/gglwe_ct.rsRaw
  1use poulpy_hal::{
  2    api::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow},
  3    layouts::{Backend, Module, Scratch, ScratchOwned},
  4    source::Source,
  5};
  6
  7use crate::{
  8    GGLWEKeyswitch, GGLWENoise, GLWESwitchingKeyEncryptSk, ScratchTakeCore,
  9    encryption::SIGMA,
 10    layouts::{
 11        GGLWEInfos, GLWESecret, GLWESecretPreparedFactory, GLWESwitchingKey, GLWESwitchingKeyLayout,
 12        GLWESwitchingKeyPreparedFactory,
 13        prepared::{GLWESecretPrepared, GLWESwitchingKeyPrepared},
 14    },
 15    noise::log2_std_noise_gglwe_product,
 16    var_noise_gglwe_product_v2,
 17};
 18
 19pub fn test_gglwe_switching_key_keyswitch<BE: Backend>(module: &Module<BE>)
 20where
 21    Module<BE>: GLWESwitchingKeyEncryptSk<BE>
 22        + GGLWEKeyswitch<BE>
 23        + GLWESwitchingKeyPreparedFactory<BE>
 24        + GLWESecretPreparedFactory<BE>
 25        + GGLWENoise<BE>,
 26    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 27    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 28{
 29    let base2k_in: usize = 17;
 30    let base2k_key: usize = 13;
 31    let base2k_out: usize = base2k_in; // MUST BE SAME
 32    let k_in: usize = 102;
 33    let max_dsize: usize = k_in.div_ceil(base2k_key);
 34
 35    for rank_in_s0s1 in 1_usize..2 {
 36        for rank_out_s0s1 in 1_usize..3 {
 37            for rank_out_s1s2 in 1_usize..3 {
 38                for dsize in 1_usize..max_dsize + 1 {
 39                    let k_ksk: usize = k_in + base2k_key * dsize;
 40                    let k_out: usize = k_ksk; // Better capture noise.
 41
 42                    let n: usize = module.n();
 43                    let dsize_in: usize = 1;
 44                    let dnum_in: usize = k_in / base2k_in;
 45                    let dnum_ksk: usize = k_in.div_ceil(base2k_key * dsize);
 46
 47                    let gglwe_s0s1_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
 48                        n: n.into(),
 49                        base2k: base2k_in.into(),
 50                        k: k_in.into(),
 51                        dnum: dnum_in.into(),
 52                        dsize: dsize_in.into(),
 53                        rank_in: rank_in_s0s1.into(),
 54                        rank_out: rank_out_s0s1.into(),
 55                    };
 56
 57                    let gglwe_s1s2_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
 58                        n: n.into(),
 59                        base2k: base2k_key.into(),
 60                        k: k_ksk.into(),
 61                        dnum: dnum_ksk.into(),
 62                        dsize: dsize.into(),
 63                        rank_in: rank_out_s0s1.into(),
 64                        rank_out: rank_out_s1s2.into(),
 65                    };
 66
 67                    let gglwe_s0s2_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
 68                        n: n.into(),
 69                        base2k: base2k_out.into(),
 70                        k: k_out.into(),
 71                        dnum: dnum_in.into(),
 72                        dsize: dsize_in.into(),
 73                        rank_in: rank_in_s0s1.into(),
 74                        rank_out: rank_out_s1s2.into(),
 75                    };
 76
 77                    let mut gglwe_s0s1: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_s0s1_infos);
 78                    let mut gglwe_s1s2: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_s1s2_infos);
 79                    let mut gglwe_s0s2: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_s0s2_infos);
 80
 81                    let mut source_xs: Source = Source::new([0u8; 32]);
 82                    let mut source_xe: Source = Source::new([0u8; 32]);
 83                    let mut source_xa: Source = Source::new([0u8; 32]);
 84
 85                    let mut scratch_enc: ScratchOwned<BE> = ScratchOwned::alloc(
 86                        GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &gglwe_s0s1_infos)
 87                            | GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &gglwe_s1s2_infos)
 88                            | GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &gglwe_s0s2_infos),
 89                    );
 90                    let mut scratch_apply: ScratchOwned<BE> = ScratchOwned::alloc(GLWESwitchingKey::keyswitch_tmp_bytes(
 91                        module,
 92                        &gglwe_s0s2_infos,
 93                        &gglwe_s0s1_infos,
 94                        &gglwe_s1s2_infos,
 95                    ));
 96
 97                    let mut sk0: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_in_s0s1.into());
 98                    sk0.fill_ternary_prob(0.5, &mut source_xs);
 99
100                    let mut sk1: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out_s0s1.into());
101                    sk1.fill_ternary_prob(0.5, &mut source_xs);
102
103                    let mut sk2: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out_s1s2.into());
104                    sk2.fill_ternary_prob(0.5, &mut source_xs);
105
106                    let mut sk2_prepared: GLWESecretPrepared<Vec<u8>, BE> =
107                        GLWESecretPrepared::alloc(module, rank_out_s1s2.into());
108                    sk2_prepared.prepare(module, &sk2);
109
110                    // gglwe_{s1}(s0) = s0 -> s1
111                    gglwe_s0s1.encrypt_sk(
112                        module,
113                        &sk0,
114                        &sk1,
115                        &mut source_xa,
116                        &mut source_xe,
117                        scratch_enc.borrow(),
118                    );
119
120                    // gglwe_{s2}(s1) -> s1 -> s2
121                    gglwe_s1s2.encrypt_sk(
122                        module,
123                        &sk1,
124                        &sk2,
125                        &mut source_xa,
126                        &mut source_xe,
127                        scratch_enc.borrow(),
128                    );
129
130                    let mut gglwe_s1s2_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
131                        GLWESwitchingKeyPrepared::alloc_from_infos(module, &gglwe_s1s2);
132                    gglwe_s1s2_prepared.prepare(module, &gglwe_s1s2, scratch_apply.borrow());
133
134                    // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
135                    gglwe_s0s2.keyswitch(
136                        module,
137                        &gglwe_s0s1,
138                        &gglwe_s1s2_prepared,
139                        scratch_apply.borrow(),
140                    );
141
142                    let max_noise: f64 = var_noise_gglwe_product_v2(
143                        module.n() as f64,
144                        k_ksk,
145                        dnum_ksk,
146                        dsize,
147                        base2k_key,
148                        0.5,
149                        0.5,
150                        0f64,
151                        SIGMA * SIGMA,
152                        0f64,
153                        rank_out_s0s1 as f64,
154                    )
155                    .sqrt()
156                    .log2();
157
158                    for row in 0..gglwe_s0s2.dnum().as_usize() {
159                        for col in 0..gglwe_s0s2.rank_in().as_usize() {
160                            assert!(
161                                gglwe_s0s2
162                                    .key
163                                    .noise(
164                                        module,
165                                        row,
166                                        col,
167                                        &sk0.data,
168                                        &sk2_prepared,
169                                        scratch_apply.borrow()
170                                    )
171                                    .std()
172                                    .log2()
173                                    <= max_noise + 0.5
174                            )
175                        }
176                    }
177                }
178            }
179        }
180    }
181}
182
183#[allow(clippy::too_many_arguments)]
184pub fn test_gglwe_switching_key_keyswitch_inplace<BE: Backend>(module: &Module<BE>)
185where
186    Module<BE>: GLWESwitchingKeyEncryptSk<BE>
187        + GGLWEKeyswitch<BE>
188        + GLWESecretPreparedFactory<BE>
189        + GGLWENoise<BE>
190        + GLWESwitchingKeyPreparedFactory<BE>,
191    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
192    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
193{
194    let base2k_out: usize = 17;
195    let base2k_key: usize = 13;
196    let k_out: usize = 102;
197    let max_dsize: usize = k_out.div_ceil(base2k_key);
198
199    for rank_in in 1_usize..3 {
200        for rank_out in 1_usize..3 {
201            for dsize in 1_usize..max_dsize + 1 {
202                let k_ksk: usize = k_out + base2k_key * dsize;
203
204                let n: usize = module.n();
205                let dsize_in: usize = 1;
206
207                let dnum_in: usize = k_out / base2k_out;
208                let dnum_ksk: usize = k_out.div_ceil(base2k_key * dsize);
209
210                let gglwe_s0s1_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
211                    n: n.into(),
212                    base2k: base2k_out.into(),
213                    k: k_out.into(),
214                    dnum: dnum_in.into(),
215                    dsize: dsize_in.into(),
216                    rank_in: rank_in.into(),
217                    rank_out: rank_out.into(),
218                };
219
220                let gglwe_s1s2_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
221                    n: n.into(),
222                    base2k: base2k_key.into(),
223                    k: k_ksk.into(),
224                    dnum: dnum_ksk.into(),
225                    dsize: dsize.into(),
226                    rank_in: rank_out.into(),
227                    rank_out: rank_out.into(),
228                };
229
230                let mut gglwe_s0s1: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_s0s1_infos);
231                let mut gglwe_s1s2: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_s1s2_infos);
232
233                let mut source_xs: Source = Source::new([0u8; 32]);
234                let mut source_xe: Source = Source::new([0u8; 32]);
235                let mut source_xa: Source = Source::new([0u8; 32]);
236
237                let mut scratch_enc: ScratchOwned<BE> = ScratchOwned::alloc(
238                    GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &gglwe_s0s1_infos)
239                        | GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &gglwe_s1s2_infos),
240                );
241                let mut scratch_apply: ScratchOwned<BE> = ScratchOwned::alloc(GLWESwitchingKey::keyswitch_tmp_bytes(
242                    module,
243                    &gglwe_s0s1_infos,
244                    &gglwe_s0s1_infos,
245                    &gglwe_s1s2_infos,
246                ));
247
248                let var_xs: f64 = 0.5;
249
250                let mut sk0: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_in.into());
251                sk0.fill_ternary_prob(var_xs, &mut source_xs);
252
253                let mut sk1: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out.into());
254                sk1.fill_ternary_prob(var_xs, &mut source_xs);
255
256                let mut sk2: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out.into());
257                sk2.fill_ternary_prob(var_xs, &mut source_xs);
258
259                let mut sk2_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank_out.into());
260                sk2_prepared.prepare(module, &sk2);
261
262                // gglwe_{s1}(s0) = s0 -> s1
263                gglwe_s0s1.encrypt_sk(
264                    module,
265                    &sk0,
266                    &sk1,
267                    &mut source_xa,
268                    &mut source_xe,
269                    scratch_enc.borrow(),
270                );
271
272                // gglwe_{s2}(s1) -> s1 -> s2
273                gglwe_s1s2.encrypt_sk(
274                    module,
275                    &sk1,
276                    &sk2,
277                    &mut source_xa,
278                    &mut source_xe,
279                    scratch_enc.borrow(),
280                );
281
282                let mut gglwe_s1s2_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
283                    GLWESwitchingKeyPrepared::alloc_from_infos(module, &gglwe_s1s2);
284                gglwe_s1s2_prepared.prepare(module, &gglwe_s1s2, scratch_apply.borrow());
285
286                // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
287                gglwe_s0s1.keyswitch_inplace(module, &gglwe_s1s2_prepared, scratch_apply.borrow());
288
289                let gglwe_s0s2: GLWESwitchingKey<Vec<u8>> = gglwe_s0s1;
290
291                let max_noise: f64 = log2_std_noise_gglwe_product(
292                    n as f64,
293                    base2k_key * dsize,
294                    var_xs,
295                    var_xs,
296                    0f64,
297                    SIGMA * SIGMA,
298                    0f64,
299                    rank_out as f64,
300                    k_out,
301                    k_ksk,
302                );
303
304                for row in 0..gglwe_s0s2.dnum().as_usize() {
305                    for col in 0..gglwe_s0s2.rank_in().as_usize() {
306                        assert!(
307                            gglwe_s0s2
308                                .key
309                                .noise(
310                                    module,
311                                    row,
312                                    col,
313                                    &sk0.data,
314                                    &sk2_prepared,
315                                    scratch_apply.borrow()
316                                )
317                                .std()
318                                .log2()
319                                <= max_noise + 0.5
320                        )
321                    }
322                }
323            }
324        }
325    }
326}