1use poulpy_hal::{
  2    api::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAutomorphism, VecZnxFillUniform},
  3    layouts::{Backend, GaloisElement, Module, Scratch, ScratchOwned},
  4    source::Source,
  5};
  6
  7use crate::{
  8    GGLWEKeyswitch, GLWEAutomorphismKeyCompressedEncryptSk, GLWEAutomorphismKeyEncryptSk, GLWESwitchingKeyCompressedEncryptSk,
  9    GLWESwitchingKeyEncryptSk, ScratchTakeCore,
 10    encryption::SIGMA,
 11    layouts::{
 12        GGLWEInfos, GLWEAutomorphismKey, GLWEAutomorphismKeyDecompress, GLWEAutomorphismKeyLayout, GLWEInfos, GLWESecret,
 13        GLWESecretPreparedFactory, GLWESwitchingKeyDecompress, LWEInfos, compressed::GLWEAutomorphismKeyCompressed,
 14        prepared::GLWESecretPrepared,
 15    },
 16    noise::GGLWENoise,
 17};
 18
 19pub fn test_gglwe_automorphism_key_encrypt_sk<BE: Backend>(module: &Module<BE>)
 20where
 21    Module<BE>: GLWEAutomorphismKeyEncryptSk<BE>
 22        + GGLWEKeyswitch<BE>
 23        + GLWESecretPreparedFactory<BE>
 24        + GLWESwitchingKeyEncryptSk<BE>
 25        + GLWESwitchingKeyCompressedEncryptSk<BE>
 26        + GLWESwitchingKeyDecompress
 27        + GGLWENoise<BE>
 28        + VecZnxFillUniform
 29        + VecZnxAutomorphism,
 30    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 31    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 32{
 33    let base2k: usize = 12;
 34    let k_ksk: usize = 60;
 35    let dsize: usize = k_ksk.div_ceil(base2k) - 1;
 36    for rank in 1_usize..3 {
 37        for di in 1..dsize + 1 {
 38            let n: usize = module.n();
 39            let dnum: usize = (k_ksk - di * base2k) / (di * base2k);
 40
 41            let atk_infos: GLWEAutomorphismKeyLayout = GLWEAutomorphismKeyLayout {
 42                n: n.into(),
 43                base2k: base2k.into(),
 44                k: k_ksk.into(),
 45                dnum: dnum.into(),
 46                dsize: di.into(),
 47                rank: rank.into(),
 48            };
 49
 50            let mut atk: GLWEAutomorphismKey<Vec<u8>> = GLWEAutomorphismKey::alloc_from_infos(&atk_infos);
 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            let mut scratch: ScratchOwned<BE> =
 57                ScratchOwned::alloc(GLWEAutomorphismKey::encrypt_sk_tmp_bytes(module, &atk_infos));
 58
 59            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&atk_infos);
 60            sk.fill_ternary_prob(0.5, &mut source_xs);
 61
 62            let p = -5;
 63
 64            atk.encrypt_sk(module, p, &sk, &mut source_xa, &mut source_xe, scratch.borrow());
 65
 66            let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
 67            (0..atk.rank().into()).for_each(|i| {
 68                module.vec_znx_automorphism(
 69                    module.galois_element_inv(p),
 70                    &mut sk_out.data.as_vec_znx_mut(),
 71                    i,
 72                    &sk.data.as_vec_znx(),
 73                    i,
 74                );
 75            });
 76            let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, sk_out.rank());
 77            sk_out_prepared.prepare(module, &sk_out);
 78
 79            let max_noise: f64 = SIGMA.log2() - (atk.k().as_usize() as f64) + 0.5;
 80
 81            for row in 0..atk.dnum().as_usize() {
 82                for col in 0..atk.rank().as_usize() {
 83                    assert!(
 84                        atk.key
 85                            .noise(module, row, col, &sk.data, &sk_out_prepared, scratch.borrow())
 86                            .std()
 87                            .log2()
 88                            <= max_noise
 89                    )
 90                }
 91            }
 92        }
 93    }
 94}
 95
 96pub fn test_gglwe_automorphism_key_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
 97where
 98    Module<BE>: GLWEAutomorphismKeyCompressedEncryptSk<BE>
 99        + GGLWEKeyswitch<BE>
100        + GLWESecretPreparedFactory<BE>
101        + GLWESwitchingKeyEncryptSk<BE>
102        + GLWESwitchingKeyCompressedEncryptSk<BE>
103        + GLWEAutomorphismKeyDecompress
104        + VecZnxAutomorphism
105        + VecZnxFillUniform
106        + GGLWENoise<BE>,
107    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
108    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
109{
110    let base2k: usize = 12;
111    let k_ksk: usize = 60;
112    let max_dsize: usize = k_ksk.div_ceil(base2k) - 1;
113    for rank in 2_usize..3 {
114        for dsize in 1..max_dsize + 1 {
115            let n: usize = module.n();
116            let dnum: usize = (k_ksk - dsize * base2k) / (dsize * base2k);
117
118            let atk_infos: GLWEAutomorphismKeyLayout = GLWEAutomorphismKeyLayout {
119                n: n.into(),
120                base2k: base2k.into(),
121                k: k_ksk.into(),
122                dnum: dnum.into(),
123                dsize: dsize.into(),
124                rank: rank.into(),
125            };
126
127            let mut atk_compressed: GLWEAutomorphismKeyCompressed<Vec<u8>> =
128                GLWEAutomorphismKeyCompressed::alloc_from_infos(&atk_infos);
129
130            let mut source_xs: Source = Source::new([0u8; 32]);
131            let mut source_xe: Source = Source::new([0u8; 32]);
132
133            let mut scratch: ScratchOwned<BE> =
134                ScratchOwned::alloc(GLWEAutomorphismKeyCompressed::encrypt_sk_tmp_bytes(module, &atk_infos));
135
136            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&atk_infos);
137            sk.fill_ternary_prob(0.5, &mut source_xs);
138
139            let p: i64 = -5;
140
141            let seed_xa: [u8; 32] = [1u8; 32];
142
143            atk_compressed.encrypt_sk(module, p, &sk, seed_xa, &mut source_xe, scratch.borrow());
144
145            let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
146            (0..atk_compressed.rank().into()).for_each(|i| {
147                module.vec_znx_automorphism(
148                    module.galois_element_inv(p),
149                    &mut sk_out.data.as_vec_znx_mut(),
150                    i,
151                    &sk.data.as_vec_znx(),
152                    i,
153                );
154            });
155            let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, sk_out.rank());
156            sk_out_prepared.prepare(module, &sk_out);
157
158            let mut atk: GLWEAutomorphismKey<Vec<u8>> = GLWEAutomorphismKey::alloc_from_infos(&atk_infos);
159            atk.decompress(module, &atk_compressed);
160
161            let max_noise: f64 = SIGMA.log2() - (atk.k().as_usize() as f64) + 0.5;
162
163            for row in 0..atk.dnum().as_usize() {
164                for col in 0..atk.rank().as_usize() {
165                    let noise_have = atk
166                        .key
167                        .noise(module, row, col, &sk.data, &sk_out_prepared, scratch.borrow())
168                        .std()
169                        .log2();
170
171                    assert!(
172                        noise_have < max_noise + 0.5,
173                        "row:{row} col:{col} noise_have:{noise_have} > max_noise:{}",
174                        max_noise + 0.5
175                    );
176                }
177            }
178        }
179    }
180}