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, VecZnxDft},
  8    source::Source,
  9};
 10
 11use crate::{
 12    GGLWETensorKeyCompressedEncryptSk, ScratchTakeCore, TensorKeyEncryptSk,
 13    decryption::GLWEDecrypt,
 14    encryption::SIGMA,
 15    layouts::{
 16        Dsize, GLWEPlaintext, GLWESecret, GLWESecretPrepare, GLWESecretPreparedAlloc, TensorKey, TensorKeyCompressed,
 17        TensorKeyLayout, prepared::GLWESecretPrepared,
 18    },
 19};
 20
 21pub fn test_gglwe_tensor_key_encrypt_sk<BE: Backend>(module: &Module<BE>)
 22where
 23    Module<BE>: TensorKeyEncryptSk<BE>
 24        + GLWESecretPrepare<BE>
 25        + GLWESecretPreparedAlloc<BE>
 26        + GLWEDecrypt<BE>
 27        + VecZnxDftAlloc<BE>
 28        + VecZnxBigAlloc<BE>
 29        + VecZnxDftApply<BE>
 30        + SvpApplyDftToDft<BE>
 31        + VecZnxIdftApplyTmpA<BE>
 32        + VecZnxBigNormalize<BE>
 33        + VecZnxSubScalarInplace,
 34    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 35    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 36{
 37    let base2k: usize = 8;
 38    let k: usize = 54;
 39
 40    for rank in 1_usize..3 {
 41        let n: usize = module.n();
 42        let dnum: usize = k / base2k;
 43
 44        let tensor_key_infos = TensorKeyLayout {
 45            n: n.into(),
 46            base2k: base2k.into(),
 47            k: k.into(),
 48            dnum: dnum.into(),
 49            dsize: Dsize(1),
 50            rank: rank.into(),
 51        };
 52
 53        let mut tensor_key: TensorKey<Vec<u8>> = TensorKey::alloc_from_infos(&tensor_key_infos);
 54
 55        let mut source_xs: Source = Source::new([0u8; 32]);
 56        let mut source_xe: Source = Source::new([0u8; 32]);
 57        let mut source_xa: Source = Source::new([0u8; 32]);
 58
 59        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(TensorKey::encrypt_sk_tmp_bytes(module, &tensor_key_infos));
 60
 61        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&tensor_key_infos);
 62        sk.fill_ternary_prob(0.5, &mut source_xs);
 63        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
 64        sk_prepared.prepare(module, &sk);
 65
 66        tensor_key.encrypt_sk(
 67            module,
 68            &sk,
 69            &mut source_xa,
 70            &mut source_xe,
 71            scratch.borrow(),
 72        );
 73
 74        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&tensor_key_infos);
 75
 76        let mut sk_ij_dft = module.vec_znx_dft_alloc(1, 1);
 77        let mut sk_ij_big = module.vec_znx_big_alloc(1, 1);
 78        let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), 1_u32.into());
 79        let mut sk_dft: VecZnxDft<Vec<u8>, BE> = module.vec_znx_dft_alloc(rank, 1);
 80
 81        for i in 0..rank {
 82            module.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
 83        }
 84
 85        for i in 0..rank {
 86            for j in 0..rank {
 87                module.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_prepared.data, j, &sk_dft, i);
 88                module.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
 89                module.vec_znx_big_normalize(
 90                    base2k,
 91                    &mut sk_ij.data.as_vec_znx_mut(),
 92                    0,
 93                    base2k,
 94                    &sk_ij_big,
 95                    0,
 96                    scratch.borrow(),
 97                );
 98                for row_i in 0..dnum {
 99                    let ct = tensor_key.at(i, j).at(row_i, 0);
100
101                    ct.decrypt(module, &mut pt, &sk_prepared, scratch.borrow());
102
103                    module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, 0);
104
105                    let std_pt: f64 = pt.data.std(base2k, 0) * (k as f64).exp2();
106                    assert!((SIGMA - std_pt).abs() <= 0.5, "{SIGMA} {std_pt}");
107                }
108            }
109        }
110    }
111}
112
113pub fn test_gglwe_tensor_key_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
114where
115    Module<BE>: TensorKeyEncryptSk<BE>
116        + GLWESecretPrepare<BE>
117        + GLWESecretPreparedAlloc<BE>
118        + GGLWETensorKeyCompressedEncryptSk<BE>
119        + GLWEDecrypt<BE>
120        + VecZnxDftAlloc<BE>
121        + VecZnxBigAlloc<BE>
122        + VecZnxDftApply<BE>
123        + SvpApplyDftToDft<BE>
124        + VecZnxIdftApplyTmpA<BE>
125        + VecZnxSubScalarInplace
126        + VecZnxFillUniform
127        + VecZnxCopy
128        + VecZnxSwitchRing,
129    // + VecZnxNormalizeInplaceImpl<BE>
130    // + VecZnxNormalizeImpl<BE>,
131    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
132    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
133{
134    let base2k = 8;
135    let k = 54;
136    for rank in 1_usize..3 {
137        let n: usize = module.n();
138        let dnum: usize = k / base2k;
139
140        let tensor_key_infos: TensorKeyLayout = TensorKeyLayout {
141            n: n.into(),
142            base2k: base2k.into(),
143            k: k.into(),
144            dnum: dnum.into(),
145            dsize: Dsize(1),
146            rank: rank.into(),
147        };
148
149        let mut tensor_key_compressed: TensorKeyCompressed<Vec<u8>> = TensorKeyCompressed::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(TensorKeyCompressed::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: TensorKey<Vec<u8>> = TensorKey::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 = module.vec_znx_dft_alloc(1, 1);
174        let mut sk_ij_big = 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}