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> = ScratchOwned::alloc(GLWEAutomorphismKey::encrypt_sk_tmp_bytes(
 57                module, &atk_infos,
 58            ));
 59
 60            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&atk_infos);
 61            sk.fill_ternary_prob(0.5, &mut source_xs);
 62
 63            let p = -5;
 64
 65            atk.encrypt_sk(
 66                module,
 67                p,
 68                &sk,
 69                &mut source_xa,
 70                &mut source_xe,
 71                scratch.borrow(),
 72            );
 73
 74            let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
 75            (0..atk.rank().into()).for_each(|i| {
 76                module.vec_znx_automorphism(
 77                    module.galois_element_inv(p),
 78                    &mut sk_out.data.as_vec_znx_mut(),
 79                    i,
 80                    &sk.data.as_vec_znx(),
 81                    i,
 82                );
 83            });
 84            let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, sk_out.rank());
 85            sk_out_prepared.prepare(module, &sk_out);
 86
 87            let max_noise: f64 = SIGMA.log2() - (atk.k().as_usize() as f64) + 0.5;
 88
 89            for row in 0..atk.dnum().as_usize() {
 90                for col in 0..atk.rank().as_usize() {
 91                    assert!(
 92                        atk.key
 93                            .noise(
 94                                module,
 95                                row,
 96                                col,
 97                                &sk.data,
 98                                &sk_out_prepared,
 99                                scratch.borrow()
100                            )
101                            .std()
102                            .log2()
103                            <= max_noise
104                    )
105                }
106            }
107        }
108    }
109}
110
111pub fn test_gglwe_automorphism_key_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
112where
113    Module<BE>: GLWEAutomorphismKeyCompressedEncryptSk<BE>
114        + GGLWEKeyswitch<BE>
115        + GLWESecretPreparedFactory<BE>
116        + GLWESwitchingKeyEncryptSk<BE>
117        + GLWESwitchingKeyCompressedEncryptSk<BE>
118        + GLWEAutomorphismKeyDecompress
119        + VecZnxAutomorphism
120        + VecZnxFillUniform
121        + GGLWENoise<BE>,
122    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
123    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
124{
125    let base2k: usize = 12;
126    let k_ksk: usize = 60;
127    let max_dsize: usize = k_ksk.div_ceil(base2k) - 1;
128    for rank in 2_usize..3 {
129        for dsize in 1..max_dsize + 1 {
130            let n: usize = module.n();
131            let dnum: usize = (k_ksk - dsize * base2k) / (dsize * base2k);
132
133            let atk_infos: GLWEAutomorphismKeyLayout = GLWEAutomorphismKeyLayout {
134                n: n.into(),
135                base2k: base2k.into(),
136                k: k_ksk.into(),
137                dnum: dnum.into(),
138                dsize: dsize.into(),
139                rank: rank.into(),
140            };
141
142            let mut atk_compressed: GLWEAutomorphismKeyCompressed<Vec<u8>> =
143                GLWEAutomorphismKeyCompressed::alloc_from_infos(&atk_infos);
144
145            let mut source_xs: Source = Source::new([0u8; 32]);
146            let mut source_xe: Source = Source::new([0u8; 32]);
147
148            let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(GLWEAutomorphismKeyCompressed::encrypt_sk_tmp_bytes(
149                module, &atk_infos,
150            ));
151
152            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&atk_infos);
153            sk.fill_ternary_prob(0.5, &mut source_xs);
154
155            let p: i64 = -5;
156
157            let seed_xa: [u8; 32] = [1u8; 32];
158
159            atk_compressed.encrypt_sk(module, p, &sk, seed_xa, &mut source_xe, scratch.borrow());
160
161            let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
162            (0..atk_compressed.rank().into()).for_each(|i| {
163                module.vec_znx_automorphism(
164                    module.galois_element_inv(p),
165                    &mut sk_out.data.as_vec_znx_mut(),
166                    i,
167                    &sk.data.as_vec_znx(),
168                    i,
169                );
170            });
171            let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, sk_out.rank());
172            sk_out_prepared.prepare(module, &sk_out);
173
174            let mut atk: GLWEAutomorphismKey<Vec<u8>> = GLWEAutomorphismKey::alloc_from_infos(&atk_infos);
175            atk.decompress(module, &atk_compressed);
176
177            let max_noise: f64 = SIGMA.log2() - (atk.k().as_usize() as f64) + 0.5;
178
179            println!("rank: {rank} dsize: {dsize} dnum: {dnum}");
180            for row in 0..atk.dnum().as_usize() {
181                for col in 0..atk.rank().as_usize() {
182                    let noise_have = atk
183                        .key
184                        .noise(
185                            module,
186                            row,
187                            col,
188                            &sk.data,
189                            &sk_out_prepared,
190                            scratch.borrow(),
191                        )
192                        .std()
193                        .log2();
194
195                    assert!(
196                        noise_have < max_noise + 0.5,
197                        "row:{row} col:{col} noise_have:{noise_have} > max_noise:{}",
198                        max_noise + 0.5
199                    );
200                }
201            }
202        }
203    }
204}