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, GLWESecret, GLWESecretPreparedFactory, GLWESecretTensor, GLWESecretTensorFactory, GLWETensorKey,
 13        GLWETensorKeyCompressed, GLWETensorKeyLayout, 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        module.gglwe_assert_noise(&tensor_key, &sk_prepared, &sk_tensor.data, SIGMA + 0.5);
 71    }
 72}
 73
 74pub fn test_gglwe_tensor_key_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
 75where
 76    Module<BE>: GLWETensorKeyEncryptSk<BE>
 77        + GLWESecretPreparedFactory<BE>
 78        + GLWETensorKeyCompressedEncryptSk<BE>
 79        + GLWEDecrypt<BE>
 80        + GLWESecretTensorFactory<BE>
 81        + GGLWENoise<BE>
 82        + GGLWEDecompress,
 83    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 84    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 85{
 86    let base2k = 8;
 87    let k = 54;
 88    for rank in 1_usize..3 {
 89        let n: usize = module.n();
 90        let dnum: usize = k / base2k;
 91
 92        let tensor_key_infos: GLWETensorKeyLayout = GLWETensorKeyLayout {
 93            n: n.into(),
 94            base2k: base2k.into(),
 95            k: k.into(),
 96            dnum: dnum.into(),
 97            dsize: Dsize(1),
 98            rank: rank.into(),
 99        };
100
101        let mut tensor_key_compressed: GLWETensorKeyCompressed<Vec<u8>> =
102            GLWETensorKeyCompressed::alloc_from_infos(&tensor_key_infos);
103
104        let mut source_xs: Source = Source::new([0u8; 32]);
105        let mut source_xe: Source = Source::new([0u8; 32]);
106
107        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(GLWETensorKeyCompressed::encrypt_sk_tmp_bytes(
108            module,
109            &tensor_key_infos,
110        ));
111
112        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&tensor_key_infos);
113        sk.fill_ternary_prob(0.5, &mut source_xs);
114        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
115        sk_prepared.prepare(module, &sk);
116
117        let seed_xa: [u8; 32] = [1u8; 32];
118
119        tensor_key_compressed.encrypt_sk(module, &sk, seed_xa, &mut source_xe, scratch.borrow());
120
121        let mut tensor_key: GLWETensorKey<Vec<u8>> = GLWETensorKey::alloc_from_infos(&tensor_key_infos);
122        tensor_key.decompress(module, &tensor_key_compressed);
123
124        let mut sk_tensor: GLWESecretTensor<Vec<u8>> = GLWESecretTensor::alloc_from_infos(&sk);
125        sk_tensor.prepare(module, &sk, scratch.borrow());
126
127        module.gglwe_assert_noise(&tensor_key, &sk_prepared, &sk_tensor.data, SIGMA + 0.5);
128    }
129}