1use poulpy_hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApply, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
  4        VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc,
  5        VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxCopy, VecZnxDftAlloc, VecZnxDftAllocBytes, VecZnxDftFromVecZnx,
  6        VecZnxDftToVecZnxBigConsume, VecZnxDftToVecZnxBigTmpA, VecZnxFillUniform, VecZnxNormalize, VecZnxNormalizeInplace,
  7        VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, VecZnxSubScalarInplace, VecZnxSwithcDegree,
  8    },
  9    layouts::{Backend, Module, ScratchOwned, VecZnxDft},
 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        GGLWETensorKey, GLWEPlaintext, GLWESecret, Infos,
 21        compressed::{Decompress, GGLWETensorKeyCompressed},
 22        prepared::{GLWESecretPrepared, PrepareAlloc},
 23    },
 24};
 25
 26pub fn test_glwe_tensor_key_encrypt_sk<B>(module: &Module<B>, basek: usize, k: 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        + VecZnxBigAllocBytes
 46        + VecZnxBigAddInplace<B>
 47        + VecZnxCopy
 48        + VecZnxDftAlloc<B>
 49        + SvpApply<B>
 50        + VecZnxBigAlloc<B>
 51        + VecZnxDftToVecZnxBigTmpA<B>
 52        + VecZnxAddScalarInplace
 53        + VecZnxSwithcDegree
 54        + VecZnxSubScalarInplace,
 55    B: Backend
 56        + TakeVecZnxDftImpl<B>
 57        + TakeVecZnxBigImpl<B>
 58        + ScratchOwnedAllocImpl<B>
 59        + ScratchOwnedBorrowImpl<B>
 60        + ScratchAvailableImpl<B>
 61        + TakeScalarZnxImpl<B>
 62        + TakeVecZnxImpl<B>
 63        + VecZnxDftAllocBytesImpl<B>
 64        + VecZnxBigAllocBytesImpl<B>
 65        + TakeSvpPPolImpl<B>,
 66{
 67    let n: usize = module.n();
 68    let rows: usize = k / basek;
 69
 70    let mut tensor_key: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k, rows, 1, rank);
 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    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWETensorKey::encrypt_sk_scratch_space(
 77        module,
 78        n,
 79        basek,
 80        tensor_key.k(),
 81        rank,
 82    ));
 83
 84    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
 85    sk.fill_ternary_prob(0.5, &mut source_xs);
 86    let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
 87
 88    tensor_key.encrypt_sk(
 89        module,
 90        &sk,
 91        &mut source_xa,
 92        &mut source_xe,
 93        scratch.borrow(),
 94    );
 95
 96    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k);
 97
 98    let mut sk_ij_dft = module.vec_znx_dft_alloc(n, 1, 1);
 99    let mut sk_ij_big = module.vec_znx_big_alloc(n, 1, 1);
100    let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, 1);
101    let mut sk_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(n, rank, 1);
102
103    (0..rank).for_each(|i| {
104        module.vec_znx_dft_from_vec_znx(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
105    });
106
107    (0..rank).for_each(|i| {
108        (0..rank).for_each(|j| {
109            module.svp_apply(&mut sk_ij_dft, 0, &sk_prepared.data, j, &sk_dft, i);
110            module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
111            module.vec_znx_big_normalize(
112                basek,
113                &mut sk_ij.data.as_vec_znx_mut(),
114                0,
115                &sk_ij_big,
116                0,
117                scratch.borrow(),
118            );
119            (0..tensor_key.rank_in()).for_each(|col_i| {
120                (0..tensor_key.rows()).for_each(|row_i| {
121                    tensor_key
122                        .at(i, j)
123                        .at(row_i, col_i)
124                        .decrypt(module, &mut pt, &sk_prepared, scratch.borrow());
125
126                    module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, col_i);
127
128                    let std_pt: f64 = pt.data.std(basek, 0) * (k as f64).exp2();
129                    assert!((SIGMA - std_pt).abs() <= 0.5, "{} {}", SIGMA, std_pt);
130                });
131            });
132        })
133    })
134}
135
136pub fn test_glwe_tensor_key_compressed_encrypt_sk<B>(module: &Module<B>, basek: usize, k: usize, rank: usize)
137where
138    Module<B>: VecZnxDftAllocBytes
139        + VecZnxBigNormalize<B>
140        + VecZnxDftFromVecZnx<B>
141        + SvpApplyInplace<B>
142        + VecZnxDftToVecZnxBigConsume<B>
143        + VecZnxNormalizeTmpBytes
144        + VecZnxFillUniform
145        + VecZnxSubABInplace
146        + VecZnxAddInplace
147        + VecZnxNormalizeInplace<B>
148        + VecZnxAddNormal
149        + VecZnxNormalize<B>
150        + VecZnxSub
151        + SvpPrepare<B>
152        + SvpPPolAllocBytes
153        + SvpPPolAlloc<B>
154        + VecZnxBigAddSmallInplace<B>
155        + VecZnxBigAllocBytes
156        + VecZnxBigAddInplace<B>
157        + VecZnxCopy
158        + VecZnxDftAlloc<B>
159        + SvpApply<B>
160        + VecZnxBigAlloc<B>
161        + VecZnxDftToVecZnxBigTmpA<B>
162        + VecZnxAddScalarInplace
163        + VecZnxSwithcDegree
164        + VecZnxSubScalarInplace,
165    B: Backend
166        + TakeVecZnxDftImpl<B>
167        + TakeVecZnxBigImpl<B>
168        + ScratchOwnedAllocImpl<B>
169        + ScratchOwnedBorrowImpl<B>
170        + ScratchAvailableImpl<B>
171        + TakeScalarZnxImpl<B>
172        + TakeVecZnxImpl<B>
173        + VecZnxDftAllocBytesImpl<B>
174        + VecZnxBigAllocBytesImpl<B>
175        + TakeSvpPPolImpl<B>,
176{
177    let n: usize = module.n();
178    let rows: usize = k / basek;
179
180    let mut tensor_key_compressed: GGLWETensorKeyCompressed<Vec<u8>> =
181        GGLWETensorKeyCompressed::alloc(n, basek, k, rows, 1, rank);
182
183    let mut source_xs: Source = Source::new([0u8; 32]);
184    let mut source_xe: Source = Source::new([0u8; 32]);
185
186    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWETensorKeyCompressed::encrypt_sk_scratch_space(
187        module,
188        n,
189        basek,
190        tensor_key_compressed.k(),
191        rank,
192    ));
193
194    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
195    sk.fill_ternary_prob(0.5, &mut source_xs);
196    let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
197
198    let seed_xa: [u8; 32] = [1u8; 32];
199
200    tensor_key_compressed.encrypt_sk(module, &sk, seed_xa, &mut source_xe, scratch.borrow());
201
202    let mut tensor_key: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k, rows, 1, rank);
203    tensor_key.decompress(module, &tensor_key_compressed);
204
205    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k);
206
207    let mut sk_ij_dft = module.vec_znx_dft_alloc(n, 1, 1);
208    let mut sk_ij_big = module.vec_znx_big_alloc(n, 1, 1);
209    let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, 1);
210    let mut sk_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(n, rank, 1);
211
212    (0..rank).for_each(|i| {
213        module.vec_znx_dft_from_vec_znx(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
214    });
215
216    (0..rank).for_each(|i| {
217        (0..rank).for_each(|j| {
218            module.svp_apply(&mut sk_ij_dft, 0, &sk_prepared.data, j, &sk_dft, i);
219            module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
220            module.vec_znx_big_normalize(
221                basek,
222                &mut sk_ij.data.as_vec_znx_mut(),
223                0,
224                &sk_ij_big,
225                0,
226                scratch.borrow(),
227            );
228            (0..tensor_key.rank_in()).for_each(|col_i| {
229                (0..tensor_key.rows()).for_each(|row_i| {
230                    tensor_key
231                        .at(i, j)
232                        .at(row_i, col_i)
233                        .decrypt(module, &mut pt, &sk_prepared, scratch.borrow());
234
235                    module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, col_i);
236
237                    let std_pt: f64 = pt.data.std(basek, 0) * (k as f64).exp2();
238                    assert!((SIGMA - std_pt).abs() <= 0.5, "{} {}", SIGMA, std_pt);
239                });
240            });
241        })
242    })
243}