1use backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxBigAlloc, VecZnxCopy, VecZnxDftAlloc, VecZnxStd,
  4        VecZnxSubScalarInplace, VecZnxSwithcDegree, VmpPMatAlloc, VmpPMatPrepare,
  5    },
  6    layouts::{Backend, Module, ScratchOwned, VecZnxDft},
  7    oep::{
  8        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
  9        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
 10    },
 11};
 12use sampling::source::Source;
 13
 14use crate::{
 15    layouts::{GGLWETensorKey, GLWEPlaintext, GLWESecret, Infos, compressed::GGLWETensorKeyCompressed, prepared::GLWESecretExec},
 16    trait_families::{Decompress, GLWEDecryptFamily},
 17};
 18
 19use crate::trait_families::{GGLWEEncryptSkFamily, GGLWETensorKeyEncryptSkFamily, GLWESecretExecModuleFamily};
 20
 21pub fn test_glwe_tensor_key_encrypt_sk<B: Backend>(module: &Module<B>, basek: usize, k: usize, sigma: f64, rank: usize)
 22where
 23    Module<B>: GGLWEEncryptSkFamily<B>
 24        + GLWESecretExecModuleFamily<B>
 25        + GLWEDecryptFamily<B>
 26        + VecZnxSwithcDegree
 27        + VecZnxAddScalarInplace
 28        + VecZnxStd
 29        + VecZnxSubScalarInplace
 30        + VmpPMatAlloc<B>
 31        + VmpPMatPrepare<B>
 32        + GGLWETensorKeyEncryptSkFamily<B>
 33        + GLWEDecryptFamily<B>
 34        + VecZnxDftAlloc<B>
 35        + VecZnxBigAlloc<B>,
 36    B: TakeVecZnxDftImpl<B>
 37        + TakeVecZnxBigImpl<B>
 38        + TakeSvpPPolImpl<B>
 39        + ScratchOwnedAllocImpl<B>
 40        + ScratchOwnedBorrowImpl<B>
 41        + ScratchAvailableImpl<B>
 42        + TakeScalarZnxImpl<B>
 43        + TakeVecZnxImpl<B>
 44        + VecZnxDftAllocBytesImpl<B>
 45        + VecZnxBigAllocBytesImpl<B>
 46        + TakeSvpPPolImpl<B>,
 47{
 48    let n: usize = module.n();
 49    let rows: usize = k / basek;
 50
 51    let mut tensor_key: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k, rows, 1, rank);
 52
 53    let mut source_xs: Source = Source::new([0u8; 32]);
 54    let mut source_xe: Source = Source::new([0u8; 32]);
 55    let mut source_xa: Source = Source::new([0u8; 32]);
 56
 57    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWETensorKey::encrypt_sk_scratch_space(
 58        module,
 59        n,
 60        basek,
 61        tensor_key.k(),
 62        rank,
 63    ));
 64
 65    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
 66    sk.fill_ternary_prob(0.5, &mut source_xs);
 67    let mut sk_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk);
 68    sk_exec.prepare(module, &sk);
 69
 70    tensor_key.encrypt_sk(
 71        module,
 72        &sk,
 73        &mut source_xa,
 74        &mut source_xe,
 75        sigma,
 76        scratch.borrow(),
 77    );
 78
 79    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k);
 80
 81    let mut sk_ij_dft = module.vec_znx_dft_alloc(n, 1, 1);
 82    let mut sk_ij_big = module.vec_znx_big_alloc(n, 1, 1);
 83    let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, 1);
 84    let mut sk_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(n, rank, 1);
 85
 86    (0..rank).for_each(|i| {
 87        module.vec_znx_dft_from_vec_znx(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
 88    });
 89
 90    (0..rank).for_each(|i| {
 91        (0..rank).for_each(|j| {
 92            module.svp_apply(&mut sk_ij_dft, 0, &sk_exec.data, j, &sk_dft, i);
 93            module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
 94            module.vec_znx_big_normalize(
 95                basek,
 96                &mut sk_ij.data.as_vec_znx_mut(),
 97                0,
 98                &sk_ij_big,
 99                0,
100                scratch.borrow(),
101            );
102            (0..tensor_key.rank_in()).for_each(|col_i| {
103                (0..tensor_key.rows()).for_each(|row_i| {
104                    tensor_key
105                        .at(i, j)
106                        .at(row_i, col_i)
107                        .decrypt(module, &mut pt, &sk_exec, scratch.borrow());
108
109                    module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, col_i);
110
111                    let std_pt: f64 = module.vec_znx_std(basek, &pt.data, 0) * (k as f64).exp2();
112                    assert!((sigma - std_pt).abs() <= 0.5, "{} {}", sigma, std_pt);
113                });
114            });
115        })
116    })
117}
118
119pub fn test_glwe_tensor_key_compressed_encrypt_sk<B: Backend>(module: &Module<B>, basek: usize, k: usize, sigma: f64, rank: usize)
120where
121    Module<B>: GGLWEEncryptSkFamily<B>
122        + GLWESecretExecModuleFamily<B>
123        + GLWEDecryptFamily<B>
124        + VecZnxSwithcDegree
125        + VecZnxAddScalarInplace
126        + VecZnxStd
127        + VecZnxSubScalarInplace
128        + VmpPMatAlloc<B>
129        + VmpPMatPrepare<B>
130        + GGLWETensorKeyEncryptSkFamily<B>
131        + GLWEDecryptFamily<B>
132        + VecZnxDftAlloc<B>
133        + VecZnxBigAlloc<B>
134        + VecZnxCopy,
135    B: TakeVecZnxDftImpl<B>
136        + TakeVecZnxBigImpl<B>
137        + TakeSvpPPolImpl<B>
138        + ScratchOwnedAllocImpl<B>
139        + ScratchOwnedBorrowImpl<B>
140        + ScratchAvailableImpl<B>
141        + TakeScalarZnxImpl<B>
142        + TakeVecZnxImpl<B>
143        + VecZnxDftAllocBytesImpl<B>
144        + VecZnxBigAllocBytesImpl<B>
145        + TakeSvpPPolImpl<B>,
146{
147    let n: usize = module.n();
148    let rows: usize = k / basek;
149
150    let mut tensor_key_compressed: GGLWETensorKeyCompressed<Vec<u8>> = GGLWETensorKeyCompressed::alloc(n, basek, k, rows, 1, rank);
151
152    let mut source_xs: Source = Source::new([0u8; 32]);
153    let mut source_xe: Source = Source::new([0u8; 32]);
154
155    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWETensorKeyCompressed::encrypt_sk_scratch_space(
156        module,
157        n,
158        basek,
159        tensor_key_compressed.k(),
160        rank,
161    ));
162
163    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
164    sk.fill_ternary_prob(0.5, &mut source_xs);
165    let mut sk_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk);
166    sk_exec.prepare(module, &sk);
167
168    let seed_xa: [u8; 32] = [1u8; 32];
169
170    tensor_key_compressed.encrypt_sk(
171        module,
172        &sk,
173        seed_xa,
174        &mut source_xe,
175        sigma,
176        scratch.borrow(),
177    );
178
179    let mut tensor_key: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k, rows, 1, rank);
180    tensor_key.decompress(module, &tensor_key_compressed);
181
182    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k);
183
184    let mut sk_ij_dft = module.vec_znx_dft_alloc(n, 1, 1);
185    let mut sk_ij_big = module.vec_znx_big_alloc(n, 1, 1);
186    let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, 1);
187    let mut sk_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(n, rank, 1);
188
189    (0..rank).for_each(|i| {
190        module.vec_znx_dft_from_vec_znx(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
191    });
192
193    (0..rank).for_each(|i| {
194        (0..rank).for_each(|j| {
195            module.svp_apply(&mut sk_ij_dft, 0, &sk_exec.data, j, &sk_dft, i);
196            module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
197            module.vec_znx_big_normalize(
198                basek,
199                &mut sk_ij.data.as_vec_znx_mut(),
200                0,
201                &sk_ij_big,
202                0,
203                scratch.borrow(),
204            );
205            (0..tensor_key.rank_in()).for_each(|col_i| {
206                (0..tensor_key.rows()).for_each(|row_i| {
207                    tensor_key
208                        .at(i, j)
209                        .at(row_i, col_i)
210                        .decrypt(module, &mut pt, &sk_exec, scratch.borrow());
211
212                    module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, col_i);
213
214                    let std_pt: f64 = module.vec_znx_std(basek, &pt.data, 0) * (k as f64).exp2();
215                    assert!((sigma - std_pt).abs() <= 0.5, "{} {}", sigma, std_pt);
216                });
217            });
218        })
219    })
220}