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        GGLWELayout, GLWESecret, GLWESecretPreparedFactory, GLWESwitchingKey, GLWESwitchingKeyCompressed,
 13        GLWESwitchingKeyDecompress,
 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                ksk.key
 78                    .assert_noise(module, &sk_out_prepared, &sk_in.data, SIGMA);
 79            }
 80        }
 81    }
 82}
 83
 84pub fn test_gglwe_switching_key_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
 85where
 86    Module<BE>: GGLWEEncryptSk<BE>
 87        + GGLWEPreparedFactory<BE>
 88        + GGLWEKeyswitch<BE>
 89        + GLWEDecrypt<BE>
 90        + GLWESecretPreparedFactory<BE>
 91        + GLWESwitchingKeyEncryptSk<BE>
 92        + GLWESwitchingKeyCompressedEncryptSk<BE>
 93        + GLWESwitchingKeyDecompress
 94        + GGLWENoise<BE>
 95        + VecZnxFillUniform,
 96    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 97    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 98{
 99    let n: usize = module.n();
100    let base2k: usize = 12;
101    let k_ksk: usize = 54;
102    let dsize: usize = k_ksk / base2k;
103    for rank_in in 1_usize..3 {
104        for rank_out in 1_usize..3 {
105            for di in 1_usize..dsize + 1 {
106                let dnum: usize = (k_ksk - di * base2k) / (di * base2k);
107
108                let gglwe_infos: GGLWELayout = GGLWELayout {
109                    n: n.into(),
110                    base2k: base2k.into(),
111                    k: k_ksk.into(),
112                    dnum: dnum.into(),
113                    dsize: di.into(),
114                    rank_in: rank_in.into(),
115                    rank_out: rank_out.into(),
116                };
117
118                let mut ksk_compressed: GLWESwitchingKeyCompressed<Vec<u8>> =
119                    GLWESwitchingKeyCompressed::alloc_from_infos(&gglwe_infos);
120
121                let mut source_xs: Source = Source::new([0u8; 32]);
122                let mut source_xe: Source = Source::new([0u8; 32]);
123
124                let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(GLWESwitchingKeyCompressed::encrypt_sk_tmp_bytes(
125                    module,
126                    &gglwe_infos,
127                ));
128
129                let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_in.into());
130                sk_in.fill_ternary_prob(0.5, &mut source_xs);
131
132                let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out.into());
133                sk_out.fill_ternary_prob(0.5, &mut source_xs);
134                let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank_out.into());
135                sk_out_prepared.prepare(module, &sk_out);
136
137                let seed_xa = [1u8; 32];
138
139                ksk_compressed.encrypt_sk(
140                    module,
141                    &sk_in,
142                    &sk_out,
143                    seed_xa,
144                    &mut source_xe,
145                    scratch.borrow(),
146                );
147
148                let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_infos);
149                ksk.decompress(module, &ksk_compressed);
150
151                ksk.key
152                    .assert_noise(module, &sk_out_prepared, &sk_in.data, SIGMA);
153            }
154        }
155    }
156}