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