1use poulpy_hal::{
  2    api::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow},
  3    layouts::{Backend, Module, Scratch, ScratchOwned},
  4    source::Source,
  5};
  6
  7use crate::{
  8    ScratchTakeCore,
  9    encryption::SIGMA,
 10    layouts::{
 11        GLWE, GLWELayout, GLWEPlaintext, GLWESecret, GLWESwitchingKey, GLWESwitchingKeyLayout,
 12        prepared::{GLWESecretPrepared, GLWESwitchingKeyPrepared},
 13    },
 14    noise::log2_std_noise_gglwe_product,
 15};
 16
 17#[allow(clippy::too_many_arguments)]
 18pub fn test_glwe_keyswitch<BE: Backend>(module: &Module<BE>)
 19where
 20    Module<BE>:,
 21    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 22    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 23{
 24    let base2k: usize = 12;
 25    let k_in: usize = 45;
 26    let dsize: usize = k_in.div_ceil(base2k);
 27
 28    for rank_in in 1_usize..3 {
 29        for rank_out in 1_usize..3 {
 30            for di in 1_usize..dsize + 1 {
 31                let k_ksk: usize = k_in + base2k * di;
 32                let k_out: usize = k_ksk; // better capture noise
 33
 34                let n: usize = module.n();
 35                let dnum: usize = k_in.div_ceil(base2k * dsize);
 36
 37                let glwe_in_infos: GLWELayout = GLWELayout {
 38                    n: n.into(),
 39                    base2k: base2k.into(),
 40                    k: k_in.into(),
 41                    rank: rank_in.into(),
 42                };
 43
 44                let glwe_out_infos: GLWELayout = GLWELayout {
 45                    n: n.into(),
 46                    base2k: base2k.into(),
 47                    k: k_out.into(),
 48                    rank: rank_out.into(),
 49                };
 50
 51                let key_apply: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
 52                    n: n.into(),
 53                    base2k: base2k.into(),
 54                    k: k_ksk.into(),
 55                    dnum: dnum.into(),
 56                    dsize: di.into(),
 57                    rank_in: rank_in.into(),
 58                    rank_out: rank_out.into(),
 59                };
 60
 61                let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&key_apply);
 62                let mut glwe_in: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_in_infos);
 63                let mut glwe_out: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_out_infos);
 64                let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_in_infos);
 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                module.vec_znx_fill_uniform(base2k, &mut pt_want.data, 0, &mut source_xa);
 71
 72                let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
 73                    GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &key_apply)
 74                        | GLWE::encrypt_sk_tmp_bytes(module, &glwe_in_infos)
 75                        | GLWE::keyswitch_tmp_bytes(module, &glwe_out_infos, &glwe_in_infos, &key_apply),
 76                );
 77
 78                let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_in.into());
 79                sk_in.fill_ternary_prob(0.5, &mut source_xs);
 80                let sk_in_prepared: GLWESecretPrepared<Vec<u8>, BE> = sk_in.prepare_alloc(module, scratch.borrow());
 81
 82                let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out.into());
 83                sk_out.fill_ternary_prob(0.5, &mut source_xs);
 84                let sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = sk_out.prepare_alloc(module, scratch.borrow());
 85
 86                ksk.encrypt_sk(
 87                    module,
 88                    &sk_in,
 89                    &sk_out,
 90                    &mut source_xa,
 91                    &mut source_xe,
 92                    scratch.borrow(),
 93                );
 94
 95                glwe_in.encrypt_sk(
 96                    module,
 97                    &pt_want,
 98                    &sk_in_prepared,
 99                    &mut source_xa,
100                    &mut source_xe,
101                    scratch.borrow(),
102                );
103
104                let ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> = ksk.prepare_alloc(module, scratch.borrow());
105
106                glwe_out.keyswitch(module, &glwe_in, &ksk_prepared, scratch.borrow());
107
108                let max_noise: f64 = log2_std_noise_gglwe_product(
109                    module.n() as f64,
110                    base2k * dsize,
111                    0.5,
112                    0.5,
113                    0f64,
114                    SIGMA * SIGMA,
115                    0f64,
116                    rank_in as f64,
117                    k_in,
118                    k_ksk,
119                );
120
121                glwe_out.assert_noise(module, &sk_out_prepared, &pt_want, max_noise + 0.5);
122            }
123        }
124    }
125}
126
127pub fn test_glwe_keyswitch_inplace<BE: Backend>(module: &Module<BE>)
128where
129    Module<BE>:,
130    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
131    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
132{
133    let base2k: usize = 12;
134    let k_out: usize = 45;
135    let dsize: usize = k_out.div_ceil(base2k);
136
137    for rank in 1_usize..3 {
138        for di in 1..dsize + 1 {
139            let k_ksk: usize = k_out + base2k * di;
140
141            let n: usize = module.n();
142            let dnum: usize = k_out.div_ceil(base2k * dsize);
143
144            let glwe_out_infos: GLWELayout = GLWELayout {
145                n: n.into(),
146                base2k: base2k.into(),
147                k: k_out.into(),
148                rank: rank.into(),
149            };
150
151            let key_apply_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
152                n: n.into(),
153                base2k: base2k.into(),
154                k: k_ksk.into(),
155                dnum: dnum.into(),
156                dsize: di.into(),
157                rank_in: rank.into(),
158                rank_out: rank.into(),
159            };
160
161            let mut key_apply: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&key_apply_infos);
162            let mut glwe_out: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_out_infos);
163            let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
164
165            let mut source_xs: Source = Source::new([0u8; 32]);
166            let mut source_xe: Source = Source::new([0u8; 32]);
167            let mut source_xa: Source = Source::new([0u8; 32]);
168
169            module.vec_znx_fill_uniform(base2k, &mut pt_want.data, 0, &mut source_xa);
170
171            let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
172                GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &key_apply_infos)
173                    | GLWE::encrypt_sk_tmp_bytes(module, &glwe_out_infos)
174                    | GLWE::keyswitch_inplace_tmp_bytes(module, &glwe_out_infos, &key_apply_infos),
175            );
176
177            let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
178            sk_in.fill_ternary_prob(0.5, &mut source_xs);
179            let sk_in_prepared: GLWESecretPrepared<Vec<u8>, BE> = sk_in.prepare_alloc(module, scratch.borrow());
180
181            let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
182            sk_out.fill_ternary_prob(0.5, &mut source_xs);
183            let sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = sk_out.prepare_alloc(module, scratch.borrow());
184
185            key_apply.encrypt_sk(
186                module,
187                &sk_in,
188                &sk_out,
189                &mut source_xa,
190                &mut source_xe,
191                scratch.borrow(),
192            );
193
194            glwe_out.encrypt_sk(
195                module,
196                &pt_want,
197                &sk_in_prepared,
198                &mut source_xa,
199                &mut source_xe,
200                scratch.borrow(),
201            );
202
203            let ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> = key_apply.prepare_alloc(module, scratch.borrow());
204
205            glwe_out.keyswitch_inplace(module, &ksk_prepared, scratch.borrow());
206
207            let max_noise: f64 = log2_std_noise_gglwe_product(
208                module.n() as f64,
209                base2k * dsize,
210                0.5,
211                0.5,
212                0f64,
213                SIGMA * SIGMA,
214                0f64,
215                rank as f64,
216                k_out,
217                k_ksk,
218            );
219
220            glwe_out.assert_noise(module, &sk_out_prepared, &pt_want, max_noise + 0.5);
221        }
222    }
223}