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, GLWESecretPrepare, GLWESecretPreparedAlloc, GLWESwitchingKey, GLWESwitchingKeyCompressed,
 13        GLWESwitchingKeyDecompress,
 14        prepared::{GGLWEPrepare, GGLWEPreparedAlloc, 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        + GGLWEPrepare<BE>
 23        + GGLWEPreparedAlloc<BE>
 24        + GGLWEKeyswitch<BE>
 25        + GLWEDecrypt<BE>
 26        + GLWESecretPreparedAlloc<BE>
 27        + GLWESecretPrepare<BE>
 28        + GLWESwitchingKeyEncryptSk<BE>
 29        + VecZnxFillUniform
 30        + GGLWENoise<BE>,
 31    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 32    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 33{
 34    let n: usize = module.n();
 35    let base2k: usize = 12;
 36    let k_ksk: usize = 54;
 37    let dsize: usize = k_ksk / base2k;
 38    for rank_in in 1_usize..3 {
 39        for rank_out in 1_usize..3 {
 40            for di in 1_usize..dsize + 1 {
 41                let dnum: usize = (k_ksk - di * base2k) / (di * base2k);
 42
 43                let gglwe_infos: GGLWELayout = GGLWELayout {
 44                    n: n.into(),
 45                    base2k: base2k.into(),
 46                    k: k_ksk.into(),
 47                    dnum: dnum.into(),
 48                    dsize: di.into(),
 49                    rank_in: rank_in.into(),
 50                    rank_out: rank_out.into(),
 51                };
 52
 53                let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_infos);
 54
 55                let mut source_xs: Source = Source::new([0u8; 32]);
 56                let mut source_xe: Source = Source::new([0u8; 32]);
 57                let mut source_xa: Source = Source::new([0u8; 32]);
 58
 59                let mut scratch: ScratchOwned<BE> =
 60                    ScratchOwned::alloc(GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &gglwe_infos));
 61
 62                let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_in.into());
 63                sk_in.fill_ternary_prob(0.5, &mut source_xs);
 64
 65                let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out.into());
 66                sk_out.fill_ternary_prob(0.5, &mut source_xs);
 67                let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank_out.into());
 68                sk_out_prepared.prepare(module, &sk_out);
 69
 70                ksk.encrypt_sk(
 71                    module,
 72                    &sk_in,
 73                    &sk_out,
 74                    &mut source_xa,
 75                    &mut source_xe,
 76                    scratch.borrow(),
 77                );
 78
 79                ksk.key
 80                    .assert_noise(module, &sk_out_prepared, &sk_in.data, SIGMA);
 81            }
 82        }
 83    }
 84}
 85
 86pub fn test_gglwe_switching_key_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
 87where
 88    Module<BE>: GGLWEEncryptSk<BE>
 89        + GGLWEPrepare<BE>
 90        + GGLWEPreparedAlloc<BE>
 91        + GGLWEKeyswitch<BE>
 92        + GLWEDecrypt<BE>
 93        + GLWESecretPreparedAlloc<BE>
 94        + GLWESecretPrepare<BE>
 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 n: usize = module.n();
104    let base2k: usize = 12;
105    let k_ksk: usize = 54;
106    let dsize: usize = k_ksk / base2k;
107    for rank_in in 1_usize..3 {
108        for rank_out in 1_usize..3 {
109            for di in 1_usize..dsize + 1 {
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(&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(n.into(), 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(n.into(), 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(&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}