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