1use poulpy_hal::{
  2    api::{
  3        ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDft, VecZnxBigAlloc, VecZnxBigNormalize, VecZnxCopy, VecZnxDftAlloc, VecZnxDftApply, VecZnxFillUniform, VecZnxIdftApplyTmpA, VecZnxSubScalarInplace, VecZnxSwitchRing
  4    }, layouts::{Backend, Module, Scratch, ScratchOwned, VecZnxDft}, oep::{VecZnxNormalizeImpl, VecZnxNormalizeInplaceImpl}, source::Source
  5};
  6
  7use crate::{
  8    decryption::GLWEDecrypt, encryption::SIGMA, layouts::{
  9        prepared::GLWESecretPrepared, Dsize, GLWEPlaintext, GLWESecret, GLWESecretPrepare, GLWESecretPreparedAlloc, TensorKey, TensorKeyAlloc, TensorKeyCompressed, TensorKeyLayout
 10    }, GGLWETensorKeyCompressedEncryptSk, ScratchTakeCore, TensorKeyEncryptSk
 11};
 12
 13pub fn test_gglwe_tensor_key_encrypt_sk<BE: Backend>(module: &Module<BE>)
 14where
 15    Module<BE>: TensorKeyEncryptSk<BE>
 16        + TensorKeyAlloc
 17        + GLWESecretPrepare<BE>
 18        + GLWESecretPreparedAlloc<BE>
 19        + GLWEDecrypt<BE>
 20        + VecZnxDftAlloc<BE>
 21        + VecZnxBigAlloc<BE>
 22        + VecZnxDftApply<BE>
 23        + SvpApplyDftToDft<BE>
 24        + VecZnxIdftApplyTmpA<BE>
 25        + VecZnxBigNormalize<BE>
 26        + VecZnxSubScalarInplace,
 27    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 28    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 29{
 30    let base2k: usize = 8;
 31    let k: usize = 54;
 32
 33    for rank in 1_usize..3 {
 34        let n: usize = module.n();
 35        let dnum: usize = k / base2k;
 36
 37        let tensor_key_infos = TensorKeyLayout {
 38            n: n.into(),
 39            base2k: base2k.into(),
 40            k: k.into(),
 41            dnum: dnum.into(),
 42            dsize: Dsize(1),
 43            rank: rank.into(),
 44        };
 45
 46        let mut tensor_key: TensorKey<Vec<u8>> = TensorKey::alloc_from_infos(module, &tensor_key_infos);
 47
 48        let mut source_xs: Source = Source::new([0u8; 32]);
 49        let mut source_xe: Source = Source::new([0u8; 32]);
 50        let mut source_xa: Source = Source::new([0u8; 32]);
 51
 52        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(TensorKey::encrypt_sk_tmp_bytes(
 53            module,
 54            &tensor_key_infos,
 55        ));
 56
 57        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(module, &tensor_key_infos);
 58        sk.fill_ternary_prob(0.5, &mut source_xs);
 59        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
 60        sk_prepared.prepare(module, &sk);
 61
 62        tensor_key.encrypt_sk(
 63            module,
 64            &sk,
 65            &mut source_xa,
 66            &mut source_xe,
 67            scratch.borrow(),
 68        );
 69
 70        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(module, &tensor_key_infos);
 71
 72        let mut sk_ij_dft = module.vec_znx_dft_alloc(1, 1);
 73        let mut sk_ij_big = module.vec_znx_big_alloc(1, 1);
 74        let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, 1_u32.into());
 75        let mut sk_dft: VecZnxDft<Vec<u8>, BE> = module.vec_znx_dft_alloc(rank, 1);
 76
 77        for i in 0..rank {
 78            module.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
 79        }
 80
 81        for i in 0..rank {
 82            for j in 0..rank {
 83                module.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_prepared.data, j, &sk_dft, i);
 84                module.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
 85                module.vec_znx_big_normalize(
 86                    base2k,
 87                    &mut sk_ij.data.as_vec_znx_mut(),
 88                    0,
 89                    base2k,
 90                    &sk_ij_big,
 91                    0,
 92                    scratch.borrow(),
 93                );
 94                for row_i in 0..dnum {
 95                    let ct = tensor_key
 96                    .at(i, j)
 97                    .at(row_i, 0);
 98
 99                    ct.decrypt(module, &mut pt, &sk_prepared, scratch.borrow());
100
101                    module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, 0);
102
103                    let std_pt: f64 = pt.data.std(base2k, 0) * (k as f64).exp2();
104                    assert!((SIGMA - std_pt).abs() <= 0.5, "{SIGMA} {std_pt}");
105                }
106            }
107        }
108    }
109}
110
111pub fn test_gglwe_tensor_key_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
112where
113    Module<BE>: TensorKeyEncryptSk<BE>
114    + TensorKeyAlloc
115    + GLWESecretPrepare<BE>
116    + GLWESecretPreparedAlloc<BE>
117    + GGLWETensorKeyCompressedEncryptSk<BE>
118    + GLWEDecrypt<BE>
119    + VecZnxDftAlloc<BE>
120    + VecZnxBigAlloc<BE>
121    + VecZnxDftApply<BE>
122    + SvpApplyDftToDft<BE>
123    + VecZnxIdftApplyTmpA<BE>
124    + VecZnxSubScalarInplace
125    + VecZnxFillUniform
126    + VecZnxCopy
127    + VecZnxSwitchRing,
128    // + VecZnxNormalizeInplaceImpl<BE>
129    // + VecZnxNormalizeImpl<BE>,
130    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
131    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
132{
133    let base2k = 8;
134    let k = 54;
135    for rank in 1_usize..3 {
136        let n: usize = module.n();
137        let dnum: usize = k / base2k;
138
139        let tensor_key_infos: TensorKeyLayout = TensorKeyLayout {
140            n: n.into(),
141            base2k: base2k.into(),
142            k: k.into(),
143            dnum: dnum.into(),
144            dsize: Dsize(1),
145            rank: rank.into(),
146        };
147
148        let mut tensor_key_compressed: TensorKeyCompressed<Vec<u8>> = TensorKeyCompressed::alloc_from_infos(module, &tensor_key_infos);
149
150        let mut source_xs: Source = Source::new([0u8; 32]);
151        let mut source_xe: Source = Source::new([0u8; 32]);
152
153        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(TensorKeyCompressed::encrypt_sk_tmp_bytes(
154            module,
155            &tensor_key_infos,
156        ));
157
158        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(module, &tensor_key_infos);
159        sk.fill_ternary_prob(0.5, &mut source_xs);
160        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
161        sk_prepared.prepare(module, &sk);
162
163        let seed_xa: [u8; 32] = [1u8; 32];
164
165        tensor_key_compressed.encrypt_sk(module, &sk, seed_xa, &mut source_xe, scratch.borrow());
166
167        let mut tensor_key: TensorKey<Vec<u8>> = TensorKey::alloc_from_infos(module, &tensor_key_infos);
168        tensor_key.decompress(module, &tensor_key_compressed);
169
170        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(module, &tensor_key_infos);
171
172        let mut sk_ij_dft = module.vec_znx_dft_alloc(1, 1);
173        let mut sk_ij_big = module.vec_znx_big_alloc(1, 1);
174        let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, 1_u32.into());
175        let mut sk_dft: VecZnxDft<Vec<u8>, BE> = module.vec_znx_dft_alloc(rank, 1);
176
177        for i in 0..rank {
178            module.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
179        }
180
181        for i in 0..rank {
182            for j in 0..rank {
183                module.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_prepared.data, j, &sk_dft, i);
184                module.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
185                module.vec_znx_big_normalize(
186                    base2k,
187                    &mut sk_ij.data.as_vec_znx_mut(),
188                    0,
189                    base2k,
190                    &sk_ij_big,
191                    0,
192                    scratch.borrow(),
193                );
194                for row_i in 0..dnum {
195                    tensor_key
196                        .at(i, j)
197                        .at(row_i, 0)
198                        .decrypt(module, &mut pt, &sk_prepared, scratch.borrow());
199
200                    module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, 0);
201
202                    let std_pt: f64 = pt.data.std(base2k, 0) * (k as f64).exp2();
203                    assert!((SIGMA - std_pt).abs() <= 0.5, "{SIGMA} {std_pt}");
204                }
205            }
206        }
207    }
208}