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