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(module, &tensor_key_infos));
 50
 51        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&tensor_key_infos);
 52        sk.fill_ternary_prob(0.5, &mut source_xs);
 53        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
 54        sk_prepared.prepare(module, &sk);
 55
 56        tensor_key.encrypt_sk(module, &sk, &mut source_xa, &mut source_xe, scratch.borrow());
 57
 58        let mut sk_tensor: GLWESecretTensor<Vec<u8>> = GLWESecretTensor::alloc_from_infos(&sk);
 59        sk_tensor.prepare(module, &sk, scratch.borrow());
 60
 61        let max_noise: f64 = SIGMA.log2() - (tensor_key.k().as_usize() as f64) + 0.5;
 62
 63        for row in 0..tensor_key.dnum().as_usize() {
 64            for col in 0..tensor_key.rank_in().as_usize() {
 65                assert!(
 66                    tensor_key
 67                        .0
 68                        .noise(module, row, col, &sk_tensor.data, &sk_prepared, scratch.borrow())
 69                        .std()
 70                        .log2()
 71                        <= max_noise
 72                )
 73            }
 74        }
 75    }
 76}
 77
 78pub fn test_gglwe_tensor_key_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
 79where
 80    Module<BE>: GLWETensorKeyEncryptSk<BE>
 81        + GLWESecretPreparedFactory<BE>
 82        + GLWETensorKeyCompressedEncryptSk<BE>
 83        + GLWEDecrypt<BE>
 84        + GLWESecretTensorFactory<BE>
 85        + GGLWENoise<BE>
 86        + GGLWEDecompress,
 87    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 88    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 89{
 90    let base2k = 8;
 91    let k = 54;
 92    for rank in 1_usize..3 {
 93        let n: usize = module.n();
 94        let dnum: usize = k / base2k;
 95
 96        let tensor_key_infos: GLWETensorKeyLayout = GLWETensorKeyLayout {
 97            n: n.into(),
 98            base2k: base2k.into(),
 99            k: k.into(),
100            dnum: dnum.into(),
101            dsize: Dsize(1),
102            rank: rank.into(),
103        };
104
105        let mut tensor_key_compressed: GLWETensorKeyCompressed<Vec<u8>> =
106            GLWETensorKeyCompressed::alloc_from_infos(&tensor_key_infos);
107
108        let mut source_xs: Source = Source::new([0u8; 32]);
109        let mut source_xe: Source = Source::new([0u8; 32]);
110
111        let mut scratch: ScratchOwned<BE> =
112            ScratchOwned::alloc(GLWETensorKeyCompressed::encrypt_sk_tmp_bytes(module, &tensor_key_infos));
113
114        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&tensor_key_infos);
115        sk.fill_ternary_prob(0.5, &mut source_xs);
116        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
117        sk_prepared.prepare(module, &sk);
118
119        let seed_xa: [u8; 32] = [1u8; 32];
120
121        tensor_key_compressed.encrypt_sk(module, &sk, seed_xa, &mut source_xe, scratch.borrow());
122
123        let mut tensor_key: GLWETensorKey<Vec<u8>> = GLWETensorKey::alloc_from_infos(&tensor_key_infos);
124        tensor_key.decompress(module, &tensor_key_compressed);
125
126        let mut sk_tensor: GLWESecretTensor<Vec<u8>> = GLWESecretTensor::alloc_from_infos(&sk);
127        sk_tensor.prepare(module, &sk, scratch.borrow());
128
129        let max_noise: f64 = SIGMA.log2() - (tensor_key.k().as_usize() as f64) + 0.5;
130
131        for row in 0..tensor_key.dnum().as_usize() {
132            for col in 0..tensor_key.rank_in().as_usize() {
133                assert!(
134                    tensor_key
135                        .0
136                        .noise(module, row, col, &sk_tensor.data, &sk_prepared, scratch.borrow())
137                        .std()
138                        .log2()
139                        <= max_noise
140                )
141            }
142        }
143    }
144}