1use poulpy_hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare, VecZnxAddInplace,
  4        VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAllocBytes,
  5        VecZnxBigNormalize, VecZnxCopy, VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxFillUniform,
  6        VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, VecZnxSubScalarInplace,
  7        VecZnxSwithcDegree, 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::layouts::{
 18    GGLWESwitchingKey, GLWESecret,
 19    compressed::{Decompress, GGLWESwitchingKeyCompressed},
 20    prepared::{GLWESecretPrepared, PrepareAlloc},
 21};
 22
 23pub fn test_gglwe_switching_key_encrypt_sk<B>(
 24    module: &Module<B>,
 25    basek: usize,
 26    k_ksk: usize,
 27    digits: usize,
 28    rank_in: usize,
 29    rank_out: usize,
 30    sigma: f64,
 31) where
 32    Module<B>: VecZnxDftAllocBytes
 33        + VecZnxBigNormalize<B>
 34        + VecZnxDftFromVecZnx<B>
 35        + SvpApplyInplace<B>
 36        + VecZnxDftToVecZnxBigConsume<B>
 37        + VecZnxNormalizeTmpBytes
 38        + VecZnxFillUniform
 39        + VecZnxSubABInplace
 40        + VecZnxAddInplace
 41        + VecZnxNormalizeInplace<B>
 42        + VecZnxAddNormal
 43        + VecZnxNormalize<B>
 44        + VecZnxSub
 45        + SvpPrepare<B>
 46        + SvpPPolAllocBytes
 47        + SvpPPolAlloc<B>
 48        + VecZnxBigAddSmallInplace<B>
 49        + VecZnxSwithcDegree
 50        + VecZnxAddScalarInplace
 51        + VecZnxBigAllocBytes
 52        + VecZnxBigAddInplace<B>
 53        + VecZnxSubScalarInplace
 54        + VecZnxCopy
 55        + VmpPMatAlloc<B>
 56        + VmpPrepare<B>,
 57    B: Backend
 58        + TakeVecZnxDftImpl<B>
 59        + TakeVecZnxBigImpl<B>
 60        + ScratchOwnedAllocImpl<B>
 61        + ScratchOwnedBorrowImpl<B>
 62        + ScratchAvailableImpl<B>
 63        + TakeScalarZnxImpl<B>
 64        + TakeVecZnxImpl<B>
 65        + VecZnxDftAllocBytesImpl<B>
 66        + VecZnxBigAllocBytesImpl<B>
 67        + TakeSvpPPolImpl<B>,
 68{
 69    let n: usize = module.n();
 70    let rows: usize = (k_ksk - digits * basek) / (digits * basek);
 71
 72    let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank_in, rank_out);
 73
 74    let mut source_xs: Source = Source::new([0u8; 32]);
 75    let mut source_xe: Source = Source::new([0u8; 32]);
 76    let mut source_xa: Source = Source::new([0u8; 32]);
 77
 78    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKey::encrypt_sk_scratch_space(
 79        module, n, basek, k_ksk, rank_in, rank_out,
 80    ));
 81
 82    let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
 83    sk_in.fill_ternary_prob(0.5, &mut source_xs);
 84
 85    let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
 86    sk_out.fill_ternary_prob(0.5, &mut source_xs);
 87    let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
 88
 89    ksk.encrypt_sk(
 90        module,
 91        &sk_in,
 92        &sk_out,
 93        &mut source_xa,
 94        &mut source_xe,
 95        sigma,
 96        scratch.borrow(),
 97    );
 98
 99    ksk.key
100        .assert_noise(module, &sk_out_prepared, &sk_in.data, sigma);
101}
102
103pub fn test_gglwe_switching_key_compressed_encrypt_sk<B>(
104    module: &Module<B>,
105    basek: usize,
106    k_ksk: usize,
107    digits: usize,
108    rank_in: usize,
109    rank_out: usize,
110    sigma: f64,
111) where
112    Module<B>: VecZnxDftAllocBytes
113        + VecZnxBigNormalize<B>
114        + VecZnxDftFromVecZnx<B>
115        + SvpApplyInplace<B>
116        + VecZnxDftToVecZnxBigConsume<B>
117        + VecZnxNormalizeTmpBytes
118        + VecZnxFillUniform
119        + VecZnxSubABInplace
120        + VecZnxAddInplace
121        + VecZnxNormalizeInplace<B>
122        + VecZnxAddNormal
123        + VecZnxNormalize<B>
124        + VecZnxSub
125        + SvpPrepare<B>
126        + SvpPPolAllocBytes
127        + SvpPPolAlloc<B>
128        + VecZnxBigAddSmallInplace<B>
129        + VecZnxSwithcDegree
130        + VecZnxAddScalarInplace
131        + VecZnxBigAllocBytes
132        + VecZnxBigAddInplace<B>
133        + VecZnxSubScalarInplace
134        + VecZnxCopy
135        + VmpPMatAlloc<B>
136        + VmpPrepare<B>,
137    B: Backend
138        + TakeVecZnxDftImpl<B>
139        + TakeVecZnxBigImpl<B>
140        + ScratchOwnedAllocImpl<B>
141        + ScratchOwnedBorrowImpl<B>
142        + ScratchAvailableImpl<B>
143        + TakeScalarZnxImpl<B>
144        + TakeVecZnxImpl<B>
145        + VecZnxDftAllocBytesImpl<B>
146        + VecZnxBigAllocBytesImpl<B>
147        + TakeSvpPPolImpl<B>,
148{
149    let n: usize = module.n();
150    let rows: usize = (k_ksk - digits * basek) / (digits * basek);
151
152    let mut ksk_compressed: GGLWESwitchingKeyCompressed<Vec<u8>> =
153        GGLWESwitchingKeyCompressed::alloc(n, basek, k_ksk, rows, digits, rank_in, rank_out);
154
155    let mut source_xs: Source = Source::new([0u8; 32]);
156    let mut source_xe: Source = Source::new([0u8; 32]);
157
158    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKeyCompressed::encrypt_sk_scratch_space(
159        module, n, basek, k_ksk, rank_in, rank_out,
160    ));
161
162    let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
163    sk_in.fill_ternary_prob(0.5, &mut source_xs);
164
165    let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
166    sk_out.fill_ternary_prob(0.5, &mut source_xs);
167    let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
168
169    let seed_xa = [1u8; 32];
170
171    ksk_compressed.encrypt_sk(
172        module,
173        &sk_in,
174        &sk_out,
175        seed_xa,
176        &mut source_xe,
177        sigma,
178        scratch.borrow(),
179    );
180
181    let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank_in, rank_out);
182    ksk.decompress(module, &ksk_compressed);
183
184    ksk.key
185        .assert_noise(module, &sk_out_prepared, &sk_in.data, sigma);
186}