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