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