1use poulpy_hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
  4        VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace,
  5        VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxCopy, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxFillUniform,
  6        VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace,
  7        VecZnxSubScalarInplace, VecZnxSwitchRing, VmpPMatAlloc, VmpPrepare,
  8    },
  9    layouts::{Backend, Module, ScratchOwned},
 10    oep::{
 11        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
 12        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
 13    },
 14    source::Source,
 15};
 16
 17use crate::{
 18    encryption::SIGMA,
 19    layouts::{
 20        GGLWECiphertextLayout, GLWESecret, GLWESwitchingKey,
 21        compressed::{Decompress, GLWESwitchingKeyCompressed},
 22        prepared::{GLWESecretPrepared, PrepareAlloc},
 23    },
 24};
 25
 26pub fn test_gglwe_switching_key_encrypt_sk<B>(module: &Module<B>)
 27where
 28    Module<B>: VecZnxDftAllocBytes
 29        + VecZnxBigNormalize<B>
 30        + VecZnxDftApply<B>
 31        + SvpApplyDftToDftInplace<B>
 32        + VecZnxIdftApplyConsume<B>
 33        + VecZnxNormalizeTmpBytes
 34        + VecZnxFillUniform
 35        + VecZnxSubInplace
 36        + VecZnxAddInplace
 37        + VecZnxNormalizeInplace<B>
 38        + VecZnxAddNormal
 39        + VecZnxNormalize<B>
 40        + VecZnxSub
 41        + SvpPrepare<B>
 42        + SvpPPolAllocBytes
 43        + SvpPPolAlloc<B>
 44        + VecZnxBigAddSmallInplace<B>
 45        + VecZnxSwitchRing
 46        + VecZnxAddScalarInplace
 47        + VecZnxBigAllocBytes
 48        + VecZnxBigAddInplace<B>
 49        + VecZnxSubScalarInplace
 50        + VecZnxCopy
 51        + VmpPMatAlloc<B>
 52        + VmpPrepare<B>,
 53    B: Backend
 54        + TakeVecZnxDftImpl<B>
 55        + TakeVecZnxBigImpl<B>
 56        + ScratchOwnedAllocImpl<B>
 57        + ScratchOwnedBorrowImpl<B>
 58        + ScratchAvailableImpl<B>
 59        + TakeScalarZnxImpl<B>
 60        + TakeVecZnxImpl<B>
 61        + VecZnxDftAllocBytesImpl<B>
 62        + VecZnxBigAllocBytesImpl<B>
 63        + TakeSvpPPolImpl<B>,
 64{
 65    let base2k: usize = 12;
 66    let k_ksk: usize = 54;
 67    let dsize: usize = k_ksk / base2k;
 68    for rank_in in 1_usize..3 {
 69        for rank_out in 1_usize..3 {
 70            for di in 1_usize..dsize + 1 {
 71                let n: usize = module.n();
 72                let dnum: usize = (k_ksk - di * base2k) / (di * base2k);
 73
 74                let gglwe_infos: GGLWECiphertextLayout = GGLWECiphertextLayout {
 75                    n: n.into(),
 76                    base2k: base2k.into(),
 77                    k: k_ksk.into(),
 78                    dnum: dnum.into(),
 79                    dsize: di.into(),
 80                    rank_in: rank_in.into(),
 81                    rank_out: rank_out.into(),
 82                };
 83
 84                let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc(&gglwe_infos);
 85
 86                let mut source_xs: Source = Source::new([0u8; 32]);
 87                let mut source_xe: Source = Source::new([0u8; 32]);
 88                let mut source_xa: Source = Source::new([0u8; 32]);
 89
 90                let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GLWESwitchingKey::encrypt_sk_scratch_space(
 91                    module,
 92                    &gglwe_infos,
 93                ));
 94
 95                let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc_with(n.into(), rank_in.into());
 96                sk_in.fill_ternary_prob(0.5, &mut source_xs);
 97
 98                let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc_with(n.into(), rank_out.into());
 99                sk_out.fill_ternary_prob(0.5, &mut source_xs);
100                let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
101
102                ksk.encrypt_sk(
103                    module,
104                    &sk_in,
105                    &sk_out,
106                    &mut source_xa,
107                    &mut source_xe,
108                    scratch.borrow(),
109                );
110
111                ksk.key
112                    .assert_noise(module, &sk_out_prepared, &sk_in.data, SIGMA);
113            }
114        }
115    }
116}
117
118pub fn test_gglwe_switching_key_compressed_encrypt_sk<B>(module: &Module<B>)
119where
120    Module<B>: VecZnxDftAllocBytes
121        + VecZnxBigNormalize<B>
122        + VecZnxDftApply<B>
123        + SvpApplyDftToDftInplace<B>
124        + VecZnxIdftApplyConsume<B>
125        + VecZnxNormalizeTmpBytes
126        + VecZnxFillUniform
127        + VecZnxSubInplace
128        + VecZnxAddInplace
129        + VecZnxNormalizeInplace<B>
130        + VecZnxAddNormal
131        + VecZnxNormalize<B>
132        + VecZnxSub
133        + SvpPrepare<B>
134        + SvpPPolAllocBytes
135        + SvpPPolAlloc<B>
136        + VecZnxBigAddSmallInplace<B>
137        + VecZnxSwitchRing
138        + VecZnxAddScalarInplace
139        + VecZnxBigAllocBytes
140        + VecZnxBigAddInplace<B>
141        + VecZnxSubScalarInplace
142        + VecZnxCopy
143        + VmpPMatAlloc<B>
144        + VmpPrepare<B>,
145    B: Backend
146        + TakeVecZnxDftImpl<B>
147        + TakeVecZnxBigImpl<B>
148        + ScratchOwnedAllocImpl<B>
149        + ScratchOwnedBorrowImpl<B>
150        + ScratchAvailableImpl<B>
151        + TakeScalarZnxImpl<B>
152        + TakeVecZnxImpl<B>
153        + VecZnxDftAllocBytesImpl<B>
154        + VecZnxBigAllocBytesImpl<B>
155        + TakeSvpPPolImpl<B>,
156{
157    let base2k: usize = 12;
158    let k_ksk: usize = 54;
159    let dsize: usize = k_ksk / base2k;
160    for rank_in in 1_usize..3 {
161        for rank_out in 1_usize..3 {
162            for di in 1_usize..dsize + 1 {
163                let n: usize = module.n();
164                let dnum: usize = (k_ksk - di * base2k) / (di * base2k);
165
166                let gglwe_infos: GGLWECiphertextLayout = GGLWECiphertextLayout {
167                    n: n.into(),
168                    base2k: base2k.into(),
169                    k: k_ksk.into(),
170                    dnum: dnum.into(),
171                    dsize: di.into(),
172                    rank_in: rank_in.into(),
173                    rank_out: rank_out.into(),
174                };
175
176                let mut ksk_compressed: GLWESwitchingKeyCompressed<Vec<u8>> = GLWESwitchingKeyCompressed::alloc(&gglwe_infos);
177
178                let mut source_xs: Source = Source::new([0u8; 32]);
179                let mut source_xe: Source = Source::new([0u8; 32]);
180
181                let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GLWESwitchingKeyCompressed::encrypt_sk_scratch_space(
182                    module,
183                    &gglwe_infos,
184                ));
185
186                let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc_with(n.into(), rank_in.into());
187                sk_in.fill_ternary_prob(0.5, &mut source_xs);
188
189                let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc_with(n.into(), rank_out.into());
190                sk_out.fill_ternary_prob(0.5, &mut source_xs);
191                let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
192
193                let seed_xa = [1u8; 32];
194
195                ksk_compressed.encrypt_sk(
196                    module,
197                    &sk_in,
198                    &sk_out,
199                    seed_xa,
200                    &mut source_xe,
201                    scratch.borrow(),
202                );
203
204                let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc(&gglwe_infos);
205                ksk.decompress(module, &ksk_compressed);
206
207                ksk.key
208                    .assert_noise(module, &sk_out_prepared, &sk_in.data, SIGMA);
209            }
210        }
211    }
212}