1use backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxCopy, VecZnxSubABInplace, VmpPMatAlloc,
  4        VmpPMatPrepare,
  5    },
  6    layouts::{Backend, Module, ScalarZnx, ScratchOwned},
  7    oep::{
  8        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
  9        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
 10    },
 11};
 12use sampling::source::Source;
 13
 14use crate::{
 15    layouts::{
 16        GGSWCiphertext, GLWESecret,
 17        compressed::GGSWCiphertextCompressed,
 18        prepared::{GLWESecretPrepared, PrepareAlloc},
 19    },
 20    trait_families::{Decompress, GGSWAssertNoiseFamily},
 21};
 22
 23use crate::trait_families::{GGSWEncryptSkFamily, GLWESecretPreparedModuleFamily};
 24
 25pub fn test_ggsw_encrypt_sk<B: Backend>(module: &Module<B>, basek: usize, k: usize, digits: usize, rank: usize, sigma: f64)
 26where
 27    Module<B>: GLWESecretPreparedModuleFamily<B>
 28        + GGSWEncryptSkFamily<B>
 29        + GGSWAssertNoiseFamily<B>
 30        + VecZnxAddScalarInplace
 31        + VecZnxSubABInplace
 32        + VecZnxCopy
 33        + VmpPMatAlloc<B>
 34        + VmpPMatPrepare<B>,
 35    B: TakeVecZnxDftImpl<B>
 36        + TakeVecZnxBigImpl<B>
 37        + TakeSvpPPolImpl<B>
 38        + ScratchOwnedAllocImpl<B>
 39        + ScratchOwnedBorrowImpl<B>
 40        + ScratchAvailableImpl<B>
 41        + TakeScalarZnxImpl<B>
 42        + TakeVecZnxImpl<B>
 43        + VecZnxDftAllocBytesImpl<B>
 44        + VecZnxBigAllocBytesImpl<B>
 45        + TakeSvpPPolImpl<B>,
 46{
 47    let n: usize = module.n();
 48    let rows: usize = (k - digits * basek) / (digits * basek);
 49
 50    let mut ct: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k, rows, digits, rank);
 51
 52    let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
 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    pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
 59
 60    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGSWCiphertext::encrypt_sk_scratch_space(
 61        module, n, basek, k, rank,
 62    ));
 63
 64    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
 65    sk.fill_ternary_prob(0.5, &mut source_xs);
 66    let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
 67
 68    ct.encrypt_sk(
 69        module,
 70        &pt_scalar,
 71        &sk_prepared,
 72        &mut source_xa,
 73        &mut source_xe,
 74        sigma,
 75        scratch.borrow(),
 76    );
 77
 78    let noise_f = |_col_i: usize| -(k as f64) + sigma.log2() + 0.5;
 79
 80    ct.assert_noise(module, &sk_prepared, &pt_scalar, &noise_f);
 81}
 82
 83pub fn test_ggsw_compressed_encrypt_sk<B: Backend>(
 84    module: &Module<B>,
 85    basek: usize,
 86    k: usize,
 87    digits: usize,
 88    rank: usize,
 89    sigma: f64,
 90) where
 91    Module<B>: GLWESecretPreparedModuleFamily<B>
 92        + GGSWEncryptSkFamily<B>
 93        + GGSWAssertNoiseFamily<B>
 94        + VecZnxAddScalarInplace
 95        + VecZnxSubABInplace
 96        + VecZnxCopy
 97        + VmpPMatAlloc<B>
 98        + VmpPMatPrepare<B>,
 99    B: TakeVecZnxDftImpl<B>
100        + TakeVecZnxBigImpl<B>
101        + TakeSvpPPolImpl<B>
102        + ScratchOwnedAllocImpl<B>
103        + ScratchOwnedBorrowImpl<B>
104        + ScratchAvailableImpl<B>
105        + TakeScalarZnxImpl<B>
106        + TakeVecZnxImpl<B>
107        + VecZnxDftAllocBytesImpl<B>
108        + VecZnxBigAllocBytesImpl<B>
109        + TakeSvpPPolImpl<B>,
110{
111    let n: usize = module.n();
112    let rows: usize = (k - digits * basek) / (digits * basek);
113
114    let mut ct_compressed: GGSWCiphertextCompressed<Vec<u8>> = GGSWCiphertextCompressed::alloc(n, basek, k, rows, digits, rank);
115
116    let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
117
118    let mut source_xs: Source = Source::new([0u8; 32]);
119    let mut source_xe: Source = Source::new([0u8; 32]);
120
121    pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
122
123    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGSWCiphertextCompressed::encrypt_sk_scratch_space(
124        module, n, basek, k, rank,
125    ));
126
127    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
128    sk.fill_ternary_prob(0.5, &mut source_xs);
129    let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
130
131    let seed_xa: [u8; 32] = [1u8; 32];
132
133    ct_compressed.encrypt_sk(
134        module,
135        &pt_scalar,
136        &sk_prepared,
137        seed_xa,
138        &mut source_xe,
139        sigma,
140        scratch.borrow(),
141    );
142
143    let noise_f = |_col_i: usize| -(k as f64) + sigma.log2() + 0.5;
144
145    let mut ct: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k, rows, digits, rank);
146    ct.decompress(module, &ct_compressed);
147
148    ct.assert_noise(module, &sk_prepared, &pt_scalar, &noise_f);
149}