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, GLWESecretPreparedFactory, GLWETensorKeyCompressed, TensorKey, TensorKeyLayout,
 17        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        + 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 = TensorKeyLayout {
 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: TensorKey<Vec<u8>> = TensorKey::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(TensorKey::encrypt_sk_tmp_bytes(module, &tensor_key_infos));
 59
 60        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&tensor_key_infos);
 61        sk.fill_ternary_prob(0.5, &mut source_xs);
 62        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
 63        sk_prepared.prepare(module, &sk);
 64
 65        tensor_key.encrypt_sk(
 66            module,
 67            &sk,
 68            &mut source_xa,
 69            &mut source_xe,
 70            scratch.borrow(),
 71        );
 72
 73        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&tensor_key_infos);
 74
 75        let mut sk_ij_dft: VecZnxDft<Vec<u8>, BE> = module.vec_znx_dft_alloc(1, 1);
 76        let mut sk_ij_big: VecZnxBig<Vec<u8>, BE> = module.vec_znx_big_alloc(1, 1);
 77        let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), 1_u32.into());
 78        let mut sk_dft: VecZnxDft<Vec<u8>, BE> = module.vec_znx_dft_alloc(rank, 1);
 79
 80        for i in 0..rank {
 81            module.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
 82        }
 83
 84        for i in 0..rank {
 85            for j in 0..rank {
 86                module.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_prepared.data, j, &sk_dft, i);
 87                module.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
 88                module.vec_znx_big_normalize(
 89                    base2k,
 90                    &mut sk_ij.data.as_vec_znx_mut(),
 91                    0,
 92                    base2k,
 93                    &sk_ij_big,
 94                    0,
 95                    scratch.borrow(),
 96                );
 97                for row_i in 0..dnum {
 98                    let ct = tensor_key.at(i, j).at(row_i, 0);
 99
100                    ct.decrypt(module, &mut pt, &sk_prepared, scratch.borrow());
101
102                    module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, 0);
103
104                    let std_pt: f64 = pt.data.std(base2k, 0) * (k as f64).exp2();
105                    assert!((SIGMA - std_pt).abs() <= 0.5, "{SIGMA} {std_pt}");
106                }
107            }
108        }
109    }
110}
111
112pub fn test_gglwe_tensor_key_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
113where
114    Module<BE>: TensorKeyEncryptSk<BE>
115        + GLWESecretPreparedFactory<BE>
116        + GGLWETensorKeyCompressedEncryptSk<BE>
117        + GLWEDecrypt<BE>
118        + VecZnxDftAlloc<BE>
119        + VecZnxBigAlloc<BE>
120        + VecZnxDftApply<BE>
121        + SvpApplyDftToDft<BE>
122        + VecZnxIdftApplyTmpA<BE>
123        + VecZnxSubScalarInplace
124        + VecZnxFillUniform
125        + VecZnxCopy
126        + VecZnxSwitchRing,
127    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
128    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
129{
130    let base2k = 8;
131    let k = 54;
132    for rank in 1_usize..3 {
133        let n: usize = module.n();
134        let dnum: usize = k / base2k;
135
136        let tensor_key_infos: TensorKeyLayout = TensorKeyLayout {
137            n: n.into(),
138            base2k: base2k.into(),
139            k: k.into(),
140            dnum: dnum.into(),
141            dsize: Dsize(1),
142            rank: rank.into(),
143        };
144
145        let mut tensor_key_compressed: GLWETensorKeyCompressed<Vec<u8>> =
146            GLWETensorKeyCompressed::alloc_from_infos(&tensor_key_infos);
147
148        let mut source_xs: Source = Source::new([0u8; 32]);
149        let mut source_xe: Source = Source::new([0u8; 32]);
150
151        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(GLWETensorKeyCompressed::encrypt_sk_tmp_bytes(
152            module,
153            &tensor_key_infos,
154        ));
155
156        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&tensor_key_infos);
157        sk.fill_ternary_prob(0.5, &mut source_xs);
158        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
159        sk_prepared.prepare(module, &sk);
160
161        let seed_xa: [u8; 32] = [1u8; 32];
162
163        tensor_key_compressed.encrypt_sk(module, &sk, seed_xa, &mut source_xe, scratch.borrow());
164
165        let mut tensor_key: TensorKey<Vec<u8>> = TensorKey::alloc_from_infos(&tensor_key_infos);
166        tensor_key.decompress(module, &tensor_key_compressed);
167
168        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&tensor_key_infos);
169
170        let mut sk_ij_dft: VecZnxDft<Vec<u8>, BE> = module.vec_znx_dft_alloc(1, 1);
171        let mut sk_ij_big: VecZnxBig<Vec<u8>, BE> = module.vec_znx_big_alloc(1, 1);
172        let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), 1_u32.into());
173        let mut sk_dft: VecZnxDft<Vec<u8>, BE> = module.vec_znx_dft_alloc(rank, 1);
174
175        for i in 0..rank {
176            module.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
177        }
178
179        for i in 0..rank {
180            for j in 0..rank {
181                module.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_prepared.data, j, &sk_dft, i);
182                module.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
183                module.vec_znx_big_normalize(
184                    base2k,
185                    &mut sk_ij.data.as_vec_znx_mut(),
186                    0,
187                    base2k,
188                    &sk_ij_big,
189                    0,
190                    scratch.borrow(),
191                );
192                for row_i in 0..dnum {
193                    tensor_key
194                        .at(i, j)
195                        .at(row_i, 0)
196                        .decrypt(module, &mut pt, &sk_prepared, scratch.borrow());
197
198                    module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, 0);
199
200                    let std_pt: f64 = pt.data.std(base2k, 0) * (k as f64).exp2();
201                    assert!((SIGMA - std_pt).abs() <= 0.5, "{SIGMA} {std_pt}");
202                }
203            }
204        }
205    }
206}