1use poulpy_hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
  4        VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc,
  5        VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftAlloc, VecZnxDftAllocBytes,
  6        VecZnxDftApply, VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeInplace,
  7        VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace, 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, GGSWCiphertextLayout, GLWESecret,
 21        compressed::{Decompress, GGSWCiphertextCompressed},
 22        prepared::{GLWESecretPrepared, PrepareAlloc},
 23    },
 24};
 25
 26pub fn test_ggsw_encrypt_sk<B>(module: &Module<B>)
 27where
 28    Module<B>: VecZnxDftAllocBytes
 29        + VecZnxBigNormalize<B>
 30        + VecZnxDftApply<B>
 31        + SvpApplyDftToDftInplace<B>
 32        + VecZnxIdftApplyConsume<B>
 33        + VecZnxNormalizeTmpBytes
 34        + VecZnxFillUniform
 35        + VecZnxSubInplace
 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        + VecZnxIdftApplyTmpA<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 base2k: usize = 12;
 69    let k: usize = 54;
 70    let digits: usize = k / base2k;
 71    for rank in 1_usize..3 {
 72        for di in 1..digits + 1 {
 73            let n: usize = module.n();
 74            let rows: usize = (k - di * base2k) / (di * base2k);
 75
 76            let ggsw_infos: GGSWCiphertextLayout = GGSWCiphertextLayout {
 77                n: n.into(),
 78                base2k: base2k.into(),
 79                k: k.into(),
 80                rows: rows.into(),
 81                digits: di.into(),
 82                rank: rank.into(),
 83            };
 84
 85            let mut ct: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(&ggsw_infos);
 86
 87            let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
 88
 89            let mut source_xs: Source = Source::new([0u8; 32]);
 90            let mut source_xe: Source = Source::new([0u8; 32]);
 91            let mut source_xa: Source = Source::new([0u8; 32]);
 92
 93            pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
 94
 95            let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGSWCiphertext::encrypt_sk_scratch_space(
 96                module,
 97                &ggsw_infos,
 98            ));
 99
100            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(&ggsw_infos);
101            sk.fill_ternary_prob(0.5, &mut source_xs);
102            let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
103
104            ct.encrypt_sk(
105                module,
106                &pt_scalar,
107                &sk_prepared,
108                &mut source_xa,
109                &mut source_xe,
110                scratch.borrow(),
111            );
112
113            let noise_f = |_col_i: usize| -(k as f64) + SIGMA.log2() + 0.5;
114
115            ct.assert_noise(module, &sk_prepared, &pt_scalar, noise_f);
116        }
117    }
118}
119
120pub fn test_ggsw_compressed_encrypt_sk<B>(module: &Module<B>)
121where
122    Module<B>: VecZnxDftAllocBytes
123        + VecZnxBigNormalize<B>
124        + VecZnxDftApply<B>
125        + SvpApplyDftToDftInplace<B>
126        + VecZnxIdftApplyConsume<B>
127        + VecZnxNormalizeTmpBytes
128        + VecZnxFillUniform
129        + VecZnxSubInplace
130        + VecZnxAddInplace
131        + VecZnxNormalizeInplace<B>
132        + VecZnxAddNormal
133        + VecZnxNormalize<B>
134        + VecZnxSub
135        + SvpPrepare<B>
136        + SvpPPolAllocBytes
137        + SvpPPolAlloc<B>
138        + VecZnxBigAddSmallInplace<B>
139        + VecZnxAddScalarInplace
140        + VecZnxBigAllocBytes
141        + VecZnxBigAddInplace<B>
142        + VecZnxCopy
143        + VmpPMatAlloc<B>
144        + VmpPrepare<B>
145        + VecZnxBigAlloc<B>
146        + VecZnxDftAlloc<B>
147        + VecZnxBigNormalizeTmpBytes
148        + VecZnxIdftApplyTmpA<B>,
149    B: Backend
150        + TakeVecZnxDftImpl<B>
151        + TakeVecZnxBigImpl<B>
152        + ScratchOwnedAllocImpl<B>
153        + ScratchOwnedBorrowImpl<B>
154        + ScratchAvailableImpl<B>
155        + TakeScalarZnxImpl<B>
156        + TakeVecZnxImpl<B>
157        + VecZnxDftAllocBytesImpl<B>
158        + VecZnxBigAllocBytesImpl<B>
159        + TakeSvpPPolImpl<B>,
160{
161    let base2k: usize = 12;
162    let k: usize = 54;
163    let digits: usize = k / base2k;
164    for rank in 1_usize..3 {
165        for di in 1..digits + 1 {
166            let n: usize = module.n();
167            let rows: usize = (k - di * base2k) / (di * base2k);
168
169            let ggsw_infos: GGSWCiphertextLayout = GGSWCiphertextLayout {
170                n: n.into(),
171                base2k: base2k.into(),
172                k: k.into(),
173                rows: rows.into(),
174                digits: di.into(),
175                rank: rank.into(),
176            };
177
178            let mut ct_compressed: GGSWCiphertextCompressed<Vec<u8>> = GGSWCiphertextCompressed::alloc(&ggsw_infos);
179
180            let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
181
182            let mut source_xs: Source = Source::new([0u8; 32]);
183            let mut source_xe: Source = Source::new([0u8; 32]);
184
185            pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
186
187            let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGSWCiphertextCompressed::encrypt_sk_scratch_space(
188                module,
189                &ggsw_infos,
190            ));
191
192            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(&ggsw_infos);
193            sk.fill_ternary_prob(0.5, &mut source_xs);
194            let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
195
196            let seed_xa: [u8; 32] = [1u8; 32];
197
198            ct_compressed.encrypt_sk(
199                module,
200                &pt_scalar,
201                &sk_prepared,
202                seed_xa,
203                &mut source_xe,
204                scratch.borrow(),
205            );
206
207            let noise_f = |_col_i: usize| -(k as f64) + SIGMA.log2() + 0.5;
208
209            let mut ct: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(&ggsw_infos);
210            ct.decompress(module, &ct_compressed);
211
212            ct.assert_noise(module, &sk_prepared, &pt_scalar, noise_f);
213        }
214    }
215}