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::{
 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>(
 27    module: &Module<B>,
 28    basek: usize,
 29    k_ksk: usize,
 30    digits: usize,
 31    rank_in: usize,
 32    rank_out: usize,
 33) where
 34    Module<B>: VecZnxDftAllocBytes
 35        + VecZnxBigNormalize<B>
 36        + VecZnxDftFromVecZnx<B>
 37        + SvpApplyInplace<B>
 38        + VecZnxDftToVecZnxBigConsume<B>
 39        + VecZnxNormalizeTmpBytes
 40        + VecZnxFillUniform
 41        + VecZnxSubABInplace
 42        + VecZnxAddInplace
 43        + VecZnxNormalizeInplace<B>
 44        + VecZnxAddNormal
 45        + VecZnxNormalize<B>
 46        + VecZnxSub
 47        + SvpPrepare<B>
 48        + SvpPPolAllocBytes
 49        + SvpPPolAlloc<B>
 50        + VecZnxBigAddSmallInplace<B>
 51        + VecZnxSwithcDegree
 52        + VecZnxAddScalarInplace
 53        + VecZnxBigAllocBytes
 54        + VecZnxBigAddInplace<B>
 55        + VecZnxSubScalarInplace
 56        + VecZnxCopy
 57        + VmpPMatAlloc<B>
 58        + VmpPrepare<B>,
 59    B: Backend
 60        + TakeVecZnxDftImpl<B>
 61        + TakeVecZnxBigImpl<B>
 62        + ScratchOwnedAllocImpl<B>
 63        + ScratchOwnedBorrowImpl<B>
 64        + ScratchAvailableImpl<B>
 65        + TakeScalarZnxImpl<B>
 66        + TakeVecZnxImpl<B>
 67        + VecZnxDftAllocBytesImpl<B>
 68        + VecZnxBigAllocBytesImpl<B>
 69        + TakeSvpPPolImpl<B>,
 70{
 71    let n: usize = module.n();
 72    let rows: usize = (k_ksk - digits * basek) / (digits * basek);
 73
 74    let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, 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, n, 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
104pub fn test_gglwe_switching_key_compressed_encrypt_sk<B>(
105    module: &Module<B>,
106    basek: usize,
107    k_ksk: usize,
108    digits: usize,
109    rank_in: usize,
110    rank_out: usize,
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        scratch.borrow(),
178    );
179
180    let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank_in, rank_out);
181    ksk.decompress(module, &ksk_compressed);
182
183    ksk.key
184        .assert_noise(module, &sk_out_prepared, &sk_in.data, SIGMA);
185}