1use poulpy_hal::{
  2    api::{ScratchOwnedAlloc, ScratchOwnedBorrow},
  3    layouts::{Backend, Module, ScalarZnx, Scratch, ScratchOwned},
  4    source::Source,
  5};
  6
  7use crate::{
  8    GGSWCompressedEncryptSk, GGSWEncryptSk, GGSWNoise, ScratchTakeCore,
  9    encryption::SIGMA,
 10    layouts::{
 11        GGSW, GGSWDecompress, GGSWInfos, GGSWLayout, GLWEInfos, GLWESecret, GLWESecretPreparedFactory,
 12        compressed::GGSWCompressed, prepared::GLWESecretPrepared,
 13    },
 14};
 15
 16pub fn test_ggsw_encrypt_sk<BE: Backend>(module: &Module<BE>)
 17where
 18    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 19    Scratch<BE>: ScratchTakeCore<BE>,
 20    Module<BE>: GGSWEncryptSk<BE> + GLWESecretPreparedFactory<BE> + GGSWNoise<BE>,
 21{
 22    let base2k: usize = 12;
 23    let k: usize = 54;
 24    let dsize: usize = k / base2k;
 25    for rank in 1_usize..3 {
 26        for di in 1..dsize + 1 {
 27            let n: usize = module.n();
 28            let dnum: usize = (k - di * base2k) / (di * base2k);
 29
 30            let ggsw_infos: GGSWLayout = GGSWLayout {
 31                n: n.into(),
 32                base2k: base2k.into(),
 33                k: k.into(),
 34                dnum: dnum.into(),
 35                dsize: di.into(),
 36                rank: rank.into(),
 37            };
 38
 39            let mut ct: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_infos);
 40
 41            let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
 42
 43            let mut source_xs: Source = Source::new([0u8; 32]);
 44            let mut source_xe: Source = Source::new([0u8; 32]);
 45            let mut source_xa: Source = Source::new([0u8; 32]);
 46
 47            pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
 48
 49            let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(GGSW::encrypt_sk_tmp_bytes(module, &ggsw_infos));
 50
 51            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&ggsw_infos);
 52            sk.fill_ternary_prob(0.5, &mut source_xs);
 53
 54            let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
 55            sk_prepared.prepare(module, &sk);
 56
 57            ct.encrypt_sk(
 58                module,
 59                &pt_scalar,
 60                &sk_prepared,
 61                &mut source_xa,
 62                &mut source_xe,
 63                scratch.borrow(),
 64            );
 65
 66            let noise_f = |_col_i: usize| -(k as f64) + SIGMA.log2() + 0.5;
 67
 68            for row in 0..ct.dnum().as_usize() {
 69                for col in 0..ct.rank().as_usize() + 1 {
 70                    assert!(
 71                        ct.noise(module, row, col, &pt_scalar, &sk_prepared, scratch.borrow())
 72                            .std()
 73                            .log2()
 74                            <= noise_f(col)
 75                    )
 76                }
 77            }
 78        }
 79    }
 80}
 81
 82pub fn test_ggsw_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
 83where
 84    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 85    Scratch<BE>: ScratchTakeCore<BE>,
 86    Module<BE>: GGSWCompressedEncryptSk<BE> + GLWESecretPreparedFactory<BE> + GGSWNoise<BE> + GGSWDecompress,
 87{
 88    let base2k: usize = 12;
 89    let k: usize = 54;
 90    let dsize: usize = k / base2k;
 91    for rank in 1_usize..3 {
 92        for di in 1..dsize + 1 {
 93            let n: usize = module.n();
 94            let dnum: usize = (k - di * base2k) / (di * base2k);
 95
 96            let ggsw_infos: GGSWLayout = GGSWLayout {
 97                n: n.into(),
 98                base2k: base2k.into(),
 99                k: k.into(),
100                dnum: dnum.into(),
101                dsize: di.into(),
102                rank: rank.into(),
103            };
104
105            let mut ct_compressed: GGSWCompressed<Vec<u8>> = GGSWCompressed::alloc_from_infos(&ggsw_infos);
106
107            let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
108
109            let mut source_xs: Source = Source::new([0u8; 32]);
110            let mut source_xe: Source = Source::new([0u8; 32]);
111
112            pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
113
114            let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(GGSWCompressed::encrypt_sk_tmp_bytes(module, &ggsw_infos));
115
116            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&ggsw_infos);
117            sk.fill_ternary_prob(0.5, &mut source_xs);
118
119            let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
120            sk_prepared.prepare(module, &sk);
121
122            let seed_xa: [u8; 32] = [1u8; 32];
123
124            ct_compressed.encrypt_sk(
125                module,
126                &pt_scalar,
127                &sk_prepared,
128                seed_xa,
129                &mut source_xe,
130                scratch.borrow(),
131            );
132
133            let noise_f = |_col_i: usize| -(k as f64) + SIGMA.log2() + 0.5;
134
135            let mut ct: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_infos);
136            ct.decompress(module, &ct_compressed);
137
138            for row in 0..ct.dnum().as_usize() {
139                for col in 0..ct.rank().as_usize() + 1 {
140                    assert!(
141                        ct.noise(module, row, col, &pt_scalar, &sk_prepared, scratch.borrow())
142                            .std()
143                            .log2()
144                            <= noise_f(col)
145                    )
146                }
147            }
148        }
149    }
150}