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