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    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: VecZnxDft<Vec<u8>, BE> = module.vec_znx_dft_alloc(1, 1);
 77        let mut sk_ij_big: VecZnxBig<Vec<u8>, BE> = 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    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
130    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
131{
132    let base2k = 8;
133    let k = 54;
134    for rank in 1_usize..3 {
135        let n: usize = module.n();
136        let dnum: usize = k / base2k;
137
138        let tensor_key_infos: TensorKeyLayout = TensorKeyLayout {
139            n: n.into(),
140            base2k: base2k.into(),
141            k: k.into(),
142            dnum: dnum.into(),
143            dsize: Dsize(1),
144            rank: rank.into(),
145        };
146
147        let mut tensor_key_compressed: TensorKeyCompressed<Vec<u8>> = TensorKeyCompressed::alloc_from_infos(&tensor_key_infos);
148
149        let mut source_xs: Source = Source::new([0u8; 32]);
150        let mut source_xe: Source = Source::new([0u8; 32]);
151
152        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(TensorKeyCompressed::encrypt_sk_tmp_bytes(
153            module,
154            &tensor_key_infos,
155        ));
156
157        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&tensor_key_infos);
158        sk.fill_ternary_prob(0.5, &mut source_xs);
159        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
160        sk_prepared.prepare(module, &sk);
161
162        let seed_xa: [u8; 32] = [1u8; 32];
163
164        tensor_key_compressed.encrypt_sk(module, &sk, seed_xa, &mut source_xe, scratch.borrow());
165
166        let mut tensor_key: TensorKey<Vec<u8>> = TensorKey::alloc_from_infos(&tensor_key_infos);
167        tensor_key.decompress(module, &tensor_key_compressed);
168
169        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&tensor_key_infos);
170
171        let mut sk_ij_dft: VecZnxDft<Vec<u8>, BE> = module.vec_znx_dft_alloc(1, 1);
172        let mut sk_ij_big: VecZnxBig<Vec<u8>, BE> = module.vec_znx_big_alloc(1, 1);
173        let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), 1_u32.into());
174        let mut sk_dft: VecZnxDft<Vec<u8>, BE> = module.vec_znx_dft_alloc(rank, 1);
175
176        for i in 0..rank {
177            module.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
178        }
179
180        for i in 0..rank {
181            for j in 0..rank {
182                module.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_prepared.data, j, &sk_dft, i);
183                module.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
184                module.vec_znx_big_normalize(
185                    base2k,
186                    &mut sk_ij.data.as_vec_znx_mut(),
187                    0,
188                    base2k,
189                    &sk_ij_big,
190                    0,
191                    scratch.borrow(),
192                );
193                for row_i in 0..dnum {
194                    tensor_key
195                        .at(i, j)
196                        .at(row_i, 0)
197                        .decrypt(module, &mut pt, &sk_prepared, scratch.borrow());
198
199                    module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, 0);
200
201                    let std_pt: f64 = pt.data.std(base2k, 0) * (k as f64).exp2();
202                    assert!((SIGMA - std_pt).abs() <= 0.5, "{SIGMA} {std_pt}");
203                }
204            }
205        }
206    }
207}