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