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