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