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, VecZnxSubABInplace,
  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        GGLWESwitchingKey, GLWESecret,
 21        compressed::{Decompress, GGLWESwitchingKeyCompressed},
 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        + VecZnxSubABInplace
 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 basek: usize = 12;
 66    let k_ksk: usize = 54;
 67    let digits: usize = k_ksk / basek;
 68    (1..3).for_each(|rank_in| {
 69        (1..3).for_each(|rank_out| {
 70            (1..digits + 1).for_each(|di| {
 71                let n: usize = module.n();
 72                let rows: usize = (k_ksk - di * basek) / (di * basek);
 73
 74                let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, di, rank_in, rank_out);
 75
 76                let mut source_xs: Source = Source::new([0u8; 32]);
 77                let mut source_xe: Source = Source::new([0u8; 32]);
 78                let mut source_xa: Source = Source::new([0u8; 32]);
 79
 80                let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKey::encrypt_sk_scratch_space(
 81                    module, basek, k_ksk, rank_in, rank_out,
 82                ));
 83
 84                let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
 85                sk_in.fill_ternary_prob(0.5, &mut source_xs);
 86
 87                let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
 88                sk_out.fill_ternary_prob(0.5, &mut source_xs);
 89                let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
 90
 91                ksk.encrypt_sk(
 92                    module,
 93                    &sk_in,
 94                    &sk_out,
 95                    &mut source_xa,
 96                    &mut source_xe,
 97                    scratch.borrow(),
 98                );
 99
100                ksk.key
101                    .assert_noise(module, &sk_out_prepared, &sk_in.data, SIGMA);
102            });
103        });
104    });
105}
106
107pub fn test_gglwe_switching_key_compressed_encrypt_sk<B>(module: &Module<B>)
108where
109    Module<B>: VecZnxDftAllocBytes
110        + VecZnxBigNormalize<B>
111        + VecZnxDftApply<B>
112        + SvpApplyDftToDftInplace<B>
113        + VecZnxIdftApplyConsume<B>
114        + VecZnxNormalizeTmpBytes
115        + VecZnxFillUniform
116        + VecZnxSubABInplace
117        + VecZnxAddInplace
118        + VecZnxNormalizeInplace<B>
119        + VecZnxAddNormal
120        + VecZnxNormalize<B>
121        + VecZnxSub
122        + SvpPrepare<B>
123        + SvpPPolAllocBytes
124        + SvpPPolAlloc<B>
125        + VecZnxBigAddSmallInplace<B>
126        + VecZnxSwitchRing
127        + VecZnxAddScalarInplace
128        + VecZnxBigAllocBytes
129        + VecZnxBigAddInplace<B>
130        + VecZnxSubScalarInplace
131        + VecZnxCopy
132        + VmpPMatAlloc<B>
133        + VmpPrepare<B>,
134    B: Backend
135        + TakeVecZnxDftImpl<B>
136        + TakeVecZnxBigImpl<B>
137        + ScratchOwnedAllocImpl<B>
138        + ScratchOwnedBorrowImpl<B>
139        + ScratchAvailableImpl<B>
140        + TakeScalarZnxImpl<B>
141        + TakeVecZnxImpl<B>
142        + VecZnxDftAllocBytesImpl<B>
143        + VecZnxBigAllocBytesImpl<B>
144        + TakeSvpPPolImpl<B>,
145{
146    let basek: usize = 12;
147    let k_ksk: usize = 54;
148    let digits: usize = k_ksk / basek;
149    (1..3).for_each(|rank_in| {
150        (1..3).for_each(|rank_out| {
151            (1..digits + 1).for_each(|di| {
152                let n: usize = module.n();
153                let rows: usize = (k_ksk - di * basek) / (di * basek);
154
155                let mut ksk_compressed: GGLWESwitchingKeyCompressed<Vec<u8>> =
156                    GGLWESwitchingKeyCompressed::alloc(n, basek, k_ksk, rows, di, rank_in, rank_out);
157
158                let mut source_xs: Source = Source::new([0u8; 32]);
159                let mut source_xe: Source = Source::new([0u8; 32]);
160
161                let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKeyCompressed::encrypt_sk_scratch_space(
162                    module, basek, k_ksk, rank_in, rank_out,
163                ));
164
165                let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
166                sk_in.fill_ternary_prob(0.5, &mut source_xs);
167
168                let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
169                sk_out.fill_ternary_prob(0.5, &mut source_xs);
170                let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
171
172                let seed_xa = [1u8; 32];
173
174                ksk_compressed.encrypt_sk(
175                    module,
176                    &sk_in,
177                    &sk_out,
178                    seed_xa,
179                    &mut source_xe,
180                    scratch.borrow(),
181                );
182
183                let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, di, rank_in, rank_out);
184                ksk.decompress(module, &ksk_compressed);
185
186                ksk.key
187                    .assert_noise(module, &sk_out_prepared, &sk_in.data, SIGMA);
188            });
189        });
190    });
191}