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