1use poulpy_hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
  4        VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc,
  5        VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftAlloc, VecZnxDftAllocBytes,
  6        VecZnxDftApply, VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeInplace,
  7        VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace, VmpPMatAlloc, VmpPrepare,
  8    },
  9    layouts::{Backend, Module, ScalarZnx, ScratchOwned},
 10    oep::{
 11        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
 12        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
 13    },
 14    source::Source,
 15};
 16
 17use crate::{
 18    encryption::SIGMA,
 19    layouts::{
 20        GGSW, GGSWCiphertextLayout, GLWESecret,
 21        compressed::{Decompress, GGSWCompressed},
 22        prepared::{GLWESecretPrepared, PrepareAlloc},
 23    },
 24};
 25
 26pub fn test_ggsw_encrypt_sk<B: Backend>(module: &Module<B>)
 27where
 28    ScratchOwned<B>: ScratchOwnedAlloc<B>,
 29    Module<B>: SvpPrepare<B>,
 30{
 31    let base2k: usize = 12;
 32    let k: usize = 54;
 33    let dsize: usize = k / base2k;
 34    for rank in 1_usize..3 {
 35        for di in 1..dsize + 1 {
 36            let n: usize = module.n();
 37            let dnum: usize = (k - di * base2k) / (di * base2k);
 38
 39            let ggsw_infos: GGSWCiphertextLayout = GGSWCiphertextLayout {
 40                n: n.into(),
 41                base2k: base2k.into(),
 42                k: k.into(),
 43                dnum: dnum.into(),
 44                dsize: di.into(),
 45                rank: rank.into(),
 46            };
 47
 48            let mut ct: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_infos);
 49
 50            let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
 51
 52            let mut source_xs: Source = Source::new([0u8; 32]);
 53            let mut source_xe: Source = Source::new([0u8; 32]);
 54            let mut source_xa: Source = Source::new([0u8; 32]);
 55
 56            pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
 57
 58            let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGSW::encrypt_sk_scratch_space(module, &ggsw_infos));
 59
 60            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&ggsw_infos);
 61            sk.fill_ternary_prob(0.5, &mut source_xs);
 62            let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
 63
 64            ct.encrypt_sk(
 65                module,
 66                &pt_scalar,
 67                &sk_prepared,
 68                &mut source_xa,
 69                &mut source_xe,
 70                scratch.borrow(),
 71            );
 72
 73            let noise_f = |_col_i: usize| -(k as f64) + SIGMA.log2() + 0.5;
 74
 75            ct.assert_noise(module, &sk_prepared, &pt_scalar, noise_f);
 76        }
 77    }
 78}
 79
 80pub fn test_ggsw_compressed_encrypt_sk<B>(module: &Module<B>)
 81where
 82    Module<B>: VecZnxDftAllocBytes
 83        + VecZnxBigNormalize<B>
 84        + VecZnxDftApply<B>
 85        + SvpApplyDftToDftInplace<B>
 86        + VecZnxIdftApplyConsume<B>
 87        + VecZnxNormalizeTmpBytes
 88        + VecZnxFillUniform
 89        + VecZnxSubInplace
 90        + VecZnxAddInplace
 91        + VecZnxNormalizeInplace<B>
 92        + VecZnxAddNormal
 93        + VecZnxNormalize<B>
 94        + VecZnxSub
 95        + SvpPrepare<B>
 96        + SvpPPolAllocBytes
 97        + SvpPPolAlloc<B>
 98        + VecZnxBigAddSmallInplace<B>
 99        + VecZnxAddScalarInplace
100        + VecZnxBigAllocBytes
101        + VecZnxBigAddInplace<B>
102        + VecZnxCopy
103        + VmpPMatAlloc<B>
104        + VmpPrepare<B>
105        + VecZnxBigAlloc<B>
106        + VecZnxDftAlloc<B>
107        + VecZnxBigNormalizeTmpBytes
108        + VecZnxIdftApplyTmpA<B>,
109    B: Backend
110        + TakeVecZnxDftImpl<B>
111        + TakeVecZnxBigImpl<B>
112        + ScratchOwnedAllocImpl<B>
113        + ScratchOwnedBorrowImpl<B>
114        + ScratchAvailableImpl<B>
115        + TakeScalarZnxImpl<B>
116        + TakeVecZnxImpl<B>
117        + VecZnxDftAllocBytesImpl<B>
118        + VecZnxBigAllocBytesImpl<B>
119        + TakeSvpPPolImpl<B>,
120{
121    let base2k: usize = 12;
122    let k: usize = 54;
123    let dsize: usize = k / base2k;
124    for rank in 1_usize..3 {
125        for di in 1..dsize + 1 {
126            let n: usize = module.n();
127            let dnum: usize = (k - di * base2k) / (di * base2k);
128
129            let ggsw_infos: GGSWCiphertextLayout = GGSWCiphertextLayout {
130                n: n.into(),
131                base2k: base2k.into(),
132                k: k.into(),
133                dnum: dnum.into(),
134                dsize: di.into(),
135                rank: rank.into(),
136            };
137
138            let mut ct_compressed: GGSWCompressed<Vec<u8>> = GGSWCompressed::alloc_from_infos(&ggsw_infos);
139
140            let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
141
142            let mut source_xs: Source = Source::new([0u8; 32]);
143            let mut source_xe: Source = Source::new([0u8; 32]);
144
145            pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
146
147            let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGSWCompressed::encrypt_sk_scratch_space(
148                module,
149                &ggsw_infos,
150            ));
151
152            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&ggsw_infos);
153            sk.fill_ternary_prob(0.5, &mut source_xs);
154            let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
155
156            let seed_xa: [u8; 32] = [1u8; 32];
157
158            ct_compressed.encrypt_sk(
159                module,
160                &pt_scalar,
161                &sk_prepared,
162                seed_xa,
163                &mut source_xe,
164                scratch.borrow(),
165            );
166
167            let noise_f = |_col_i: usize| -(k as f64) + SIGMA.log2() + 0.5;
168
169            let mut ct: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_infos);
170            ct.decompress(module, &ct_compressed);
171
172            ct.assert_noise(module, &sk_prepared, &pt_scalar, noise_f);
173        }
174    }
175}