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