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