1use poulpy_hal::{
  2    api::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxFillUniform},
  3    layouts::{Backend, Module, Scratch, ScratchOwned},
  4    source::Source,
  5};
  6
  7use crate::{
  8    GGLWEEncryptSk, GGLWEKeyswitch, GLWESwitchingKeyCompressedEncryptSk, GLWESwitchingKeyEncryptSk, ScratchTakeCore,
  9    decryption::GLWEDecrypt,
 10    encryption::SIGMA,
 11    layouts::{
 12        GGLWEInfos, GGLWELayout, GLWESecret, GLWESecretPreparedFactory, GLWESwitchingKey, GLWESwitchingKeyCompressed,
 13        GLWESwitchingKeyDecompress, LWEInfos,
 14        prepared::{GGLWEPreparedFactory, GLWESecretPrepared},
 15    },
 16    noise::GGLWENoise,
 17};
 18
 19pub fn test_gglwe_switching_key_encrypt_sk<BE: Backend>(module: &Module<BE>)
 20where
 21    Module<BE>: GGLWEEncryptSk<BE>
 22        + GGLWEPreparedFactory<BE>
 23        + GGLWEKeyswitch<BE>
 24        + GLWEDecrypt<BE>
 25        + GLWESecretPreparedFactory<BE>
 26        + GLWESwitchingKeyEncryptSk<BE>
 27        + VecZnxFillUniform
 28        + GGLWENoise<BE>,
 29    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 30    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 31{
 32    let n: usize = module.n();
 33    let base2k: usize = 12;
 34    let k_ksk: usize = 54;
 35    let dsize: usize = k_ksk / base2k;
 36    for rank_in in 1_usize..3 {
 37        for rank_out in 1_usize..3 {
 38            for di in 1_usize..dsize + 1 {
 39                let dnum: usize = (k_ksk - di * base2k) / (di * base2k);
 40
 41                let gglwe_infos: GGLWELayout = GGLWELayout {
 42                    n: n.into(),
 43                    base2k: base2k.into(),
 44                    k: k_ksk.into(),
 45                    dnum: dnum.into(),
 46                    dsize: di.into(),
 47                    rank_in: rank_in.into(),
 48                    rank_out: rank_out.into(),
 49                };
 50
 51                let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_infos);
 52
 53                let mut source_xs: Source = Source::new([0u8; 32]);
 54                let mut source_xe: Source = Source::new([0u8; 32]);
 55                let mut source_xa: Source = Source::new([0u8; 32]);
 56
 57                let mut scratch: ScratchOwned<BE> =
 58                    ScratchOwned::alloc(GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &gglwe_infos));
 59
 60                let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_in.into());
 61                sk_in.fill_ternary_prob(0.5, &mut source_xs);
 62
 63                let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out.into());
 64                sk_out.fill_ternary_prob(0.5, &mut source_xs);
 65                let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank_out.into());
 66                sk_out_prepared.prepare(module, &sk_out);
 67
 68                ksk.encrypt_sk(
 69                    module,
 70                    &sk_in,
 71                    &sk_out,
 72                    &mut source_xa,
 73                    &mut source_xe,
 74                    scratch.borrow(),
 75                );
 76
 77                let max_noise: f64 = SIGMA.log2() - (ksk.k().as_usize() as f64) + 0.5;
 78
 79                for row in 0..ksk.dnum().as_usize() {
 80                    for col in 0..ksk.rank_in().as_usize() {
 81                        let noise_have = ksk
 82                            .key
 83                            .noise(
 84                                module,
 85                                row,
 86                                col,
 87                                &sk_in.data,
 88                                &sk_out_prepared,
 89                                scratch.borrow(),
 90                            )
 91                            .std()
 92                            .log2();
 93
 94                        assert!(
 95                            noise_have < max_noise + 0.5,
 96                            "row:{row} col:{col} noise_have:{noise_have} > max_noise:{}",
 97                            max_noise + 0.5
 98                        );
 99                    }
100                }
101            }
102        }
103    }
104}
105
106pub fn test_gglwe_switching_key_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
107where
108    Module<BE>: GGLWEEncryptSk<BE>
109        + GGLWEPreparedFactory<BE>
110        + GGLWEKeyswitch<BE>
111        + GLWEDecrypt<BE>
112        + GLWESecretPreparedFactory<BE>
113        + GLWESwitchingKeyEncryptSk<BE>
114        + GLWESwitchingKeyCompressedEncryptSk<BE>
115        + GLWESwitchingKeyDecompress
116        + GGLWENoise<BE>
117        + VecZnxFillUniform,
118    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
119    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
120{
121    let n: usize = module.n();
122    let base2k: usize = 12;
123    let k_ksk: usize = 54;
124    let dsize: usize = k_ksk / base2k;
125    for rank_in in 1_usize..3 {
126        for rank_out in 1_usize..3 {
127            for di in 1_usize..dsize + 1 {
128                let dnum: usize = (k_ksk - di * base2k) / (di * base2k);
129
130                let gglwe_infos: GGLWELayout = GGLWELayout {
131                    n: n.into(),
132                    base2k: base2k.into(),
133                    k: k_ksk.into(),
134                    dnum: dnum.into(),
135                    dsize: di.into(),
136                    rank_in: rank_in.into(),
137                    rank_out: rank_out.into(),
138                };
139
140                let mut ksk_compressed: GLWESwitchingKeyCompressed<Vec<u8>> =
141                    GLWESwitchingKeyCompressed::alloc_from_infos(&gglwe_infos);
142
143                let mut source_xs: Source = Source::new([0u8; 32]);
144                let mut source_xe: Source = Source::new([0u8; 32]);
145
146                let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(GLWESwitchingKeyCompressed::encrypt_sk_tmp_bytes(
147                    module,
148                    &gglwe_infos,
149                ));
150
151                let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_in.into());
152                sk_in.fill_ternary_prob(0.5, &mut source_xs);
153
154                let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out.into());
155                sk_out.fill_ternary_prob(0.5, &mut source_xs);
156                let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank_out.into());
157                sk_out_prepared.prepare(module, &sk_out);
158
159                let seed_xa = [1u8; 32];
160
161                ksk_compressed.encrypt_sk(
162                    module,
163                    &sk_in,
164                    &sk_out,
165                    seed_xa,
166                    &mut source_xe,
167                    scratch.borrow(),
168                );
169
170                let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_infos);
171                ksk.decompress(module, &ksk_compressed);
172
173                let max_noise: f64 = SIGMA.log2() - (ksk.k().as_usize() as f64) + 0.5;
174
175                for row in 0..ksk.dnum().as_usize() {
176                    for col in 0..ksk.rank_in().as_usize() {
177                        let noise_have = ksk
178                            .key
179                            .noise(
180                                module,
181                                row,
182                                col,
183                                &sk_in.data,
184                                &sk_out_prepared,
185                                scratch.borrow(),
186                            )
187                            .std()
188                            .log2();
189
190                        assert!(
191                            noise_have < max_noise + 0.5,
192                            "row:{row} col:{col} noise_have:{noise_have} > max_noise:{}",
193                            max_noise + 0.5
194                        );
195                    }
196                }
197            }
198        }
199    }
200}