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