1use poulpy_hal::{
  2    api::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow},
  3    layouts::{Backend, Module, Scratch, ScratchOwned},
  4    source::Source,
  5};
  6
  7use crate::{
  8    GGLWENoise, GLWETensorKeyCompressedEncryptSk, GLWETensorKeyEncryptSk, ScratchTakeCore,
  9    decryption::GLWEDecrypt,
 10    encryption::SIGMA,
 11    layouts::{
 12        Dsize, GGLWEDecompress, GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretPreparedFactory, GLWESecretTensor,
 13        GLWESecretTensorFactory, GLWETensorKey, GLWETensorKeyCompressed, GLWETensorKeyLayout, LWEInfos,
 14        prepared::GLWESecretPrepared,
 15    },
 16};
 17
 18pub fn test_gglwe_tensor_key_encrypt_sk<BE: Backend>(module: &Module<BE>)
 19where
 20    Module<BE>: GLWETensorKeyEncryptSk<BE>
 21        + GLWESecretPreparedFactory<BE>
 22        + GLWEDecrypt<BE>
 23        + GLWESecretTensorFactory<BE>
 24        + GGLWENoise<BE>,
 25    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 26    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 27{
 28    let base2k: usize = 8;
 29    let k: usize = 54;
 30
 31    for rank in 2_usize..3 {
 32        let n: usize = module.n();
 33        let dnum: usize = k / base2k;
 34
 35        let tensor_key_infos = GLWETensorKeyLayout {
 36            n: n.into(),
 37            base2k: base2k.into(),
 38            k: k.into(),
 39            dnum: dnum.into(),
 40            dsize: Dsize(1),
 41            rank: rank.into(),
 42        };
 43
 44        let mut tensor_key: GLWETensorKey<Vec<u8>> = GLWETensorKey::alloc_from_infos(&tensor_key_infos);
 45
 46        let mut source_xs: Source = Source::new([0u8; 32]);
 47        let mut source_xe: Source = Source::new([0u8; 32]);
 48        let mut source_xa: Source = Source::new([0u8; 32]);
 49
 50        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(GLWETensorKey::encrypt_sk_tmp_bytes(
 51            module,
 52            &tensor_key_infos,
 53        ));
 54
 55        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&tensor_key_infos);
 56        sk.fill_ternary_prob(0.5, &mut source_xs);
 57        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
 58        sk_prepared.prepare(module, &sk);
 59
 60        tensor_key.encrypt_sk(
 61            module,
 62            &sk,
 63            &mut source_xa,
 64            &mut source_xe,
 65            scratch.borrow(),
 66        );
 67
 68        let mut sk_tensor: GLWESecretTensor<Vec<u8>> = GLWESecretTensor::alloc_from_infos(&sk);
 69        sk_tensor.prepare(module, &sk, scratch.borrow());
 70
 71        let max_noise: f64 = SIGMA.log2() - (tensor_key.k().as_usize() as f64) + 0.5;
 72
 73        for row in 0..tensor_key.dnum().as_usize() {
 74            for col in 0..tensor_key.rank().as_usize() + 1 {
 75                assert!(
 76                    tensor_key
 77                        .0
 78                        .noise(
 79                            module,
 80                            row,
 81                            col,
 82                            &sk_tensor.data,
 83                            &sk_prepared,
 84                            scratch.borrow()
 85                        )
 86                        .std()
 87                        .log2()
 88                        <= max_noise
 89                )
 90            }
 91        }
 92    }
 93}
 94
 95pub fn test_gglwe_tensor_key_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
 96where
 97    Module<BE>: GLWETensorKeyEncryptSk<BE>
 98        + GLWESecretPreparedFactory<BE>
 99        + GLWETensorKeyCompressedEncryptSk<BE>
100        + GLWEDecrypt<BE>
101        + GLWESecretTensorFactory<BE>
102        + GGLWENoise<BE>
103        + GGLWEDecompress,
104    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
105    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
106{
107    let base2k = 8;
108    let k = 54;
109    for rank in 1_usize..3 {
110        let n: usize = module.n();
111        let dnum: usize = k / base2k;
112
113        let tensor_key_infos: GLWETensorKeyLayout = GLWETensorKeyLayout {
114            n: n.into(),
115            base2k: base2k.into(),
116            k: k.into(),
117            dnum: dnum.into(),
118            dsize: Dsize(1),
119            rank: rank.into(),
120        };
121
122        let mut tensor_key_compressed: GLWETensorKeyCompressed<Vec<u8>> =
123            GLWETensorKeyCompressed::alloc_from_infos(&tensor_key_infos);
124
125        let mut source_xs: Source = Source::new([0u8; 32]);
126        let mut source_xe: Source = Source::new([0u8; 32]);
127
128        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(GLWETensorKeyCompressed::encrypt_sk_tmp_bytes(
129            module,
130            &tensor_key_infos,
131        ));
132
133        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&tensor_key_infos);
134        sk.fill_ternary_prob(0.5, &mut source_xs);
135        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
136        sk_prepared.prepare(module, &sk);
137
138        let seed_xa: [u8; 32] = [1u8; 32];
139
140        tensor_key_compressed.encrypt_sk(module, &sk, seed_xa, &mut source_xe, scratch.borrow());
141
142        let mut tensor_key: GLWETensorKey<Vec<u8>> = GLWETensorKey::alloc_from_infos(&tensor_key_infos);
143        tensor_key.decompress(module, &tensor_key_compressed);
144
145        let mut sk_tensor: GLWESecretTensor<Vec<u8>> = GLWESecretTensor::alloc_from_infos(&sk);
146        sk_tensor.prepare(module, &sk, scratch.borrow());
147
148        let max_noise: f64 = SIGMA.log2() - (tensor_key.k().as_usize() as f64) + 0.5;
149
150        for row in 0..tensor_key.dnum().as_usize() {
151            for col in 0..tensor_key.rank().as_usize() + 1 {
152                assert!(
153                    tensor_key
154                        .0
155                        .noise(
156                            module,
157                            row,
158                            col,
159                            &sk_tensor.data,
160                            &sk_prepared,
161                            scratch.borrow()
162                        )
163                        .std()
164                        .log2()
165                        <= max_noise
166                )
167            }
168        }
169    }
170}