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