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        AutomorphismKeyDecompress, GLWEAutomorphismKey, GLWEAutomorphismKeyLayout, GLWEInfos, GLWESecret,
 13        GLWESecretPreparedFactory, GLWESwitchingKeyDecompress, 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().into());
 85            sk_out_prepared.prepare(module, &sk_out);
 86
 87            atk.key
 88                .assert_noise(module, &sk_out_prepared, &sk.data, SIGMA);
 89        }
 90    }
 91}
 92
 93pub fn test_gglwe_automorphism_key_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
 94where
 95    Module<BE>: GLWEAutomorphismKeyCompressedEncryptSk<BE>
 96        + GGLWEKeyswitch<BE>
 97        + GLWESecretPreparedFactory<BE>
 98        + GLWESwitchingKeyEncryptSk<BE>
 99        + GLWESwitchingKeyCompressedEncryptSk<BE>
100        + AutomorphismKeyDecompress
101        + VecZnxAutomorphism
102        + VecZnxFillUniform
103        + GGLWENoise<BE>,
104    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
105    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
106{
107    let base2k: usize = 12;
108    let k_ksk: usize = 60;
109    let dsize: usize = k_ksk.div_ceil(base2k) - 1;
110    for rank in 1_usize..3 {
111        for di in 1..dsize + 1 {
112            let n: usize = module.n();
113            let dnum: usize = (k_ksk - di * base2k) / (di * base2k);
114
115            let atk_infos: GLWEAutomorphismKeyLayout = GLWEAutomorphismKeyLayout {
116                n: n.into(),
117                base2k: base2k.into(),
118                k: k_ksk.into(),
119                dnum: dnum.into(),
120                dsize: di.into(),
121                rank: rank.into(),
122            };
123
124            let mut atk_compressed: GLWEAutomorphismKeyCompressed<Vec<u8>> =
125                GLWEAutomorphismKeyCompressed::alloc_from_infos(&atk_infos);
126
127            let mut source_xs: Source = Source::new([0u8; 32]);
128            let mut source_xe: Source = Source::new([0u8; 32]);
129
130            let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(GLWEAutomorphismKeyCompressed::encrypt_sk_tmp_bytes(
131                module, &atk_infos,
132            ));
133
134            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&atk_infos);
135            sk.fill_ternary_prob(0.5, &mut source_xs);
136
137            let p = -5;
138
139            let seed_xa: [u8; 32] = [1u8; 32];
140
141            atk_compressed.encrypt_sk(module, p, &sk, seed_xa, &mut source_xe, scratch.borrow());
142
143            let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
144            (0..atk_compressed.rank().into()).for_each(|i| {
145                module.vec_znx_automorphism(
146                    module.galois_element_inv(p),
147                    &mut sk_out.data.as_vec_znx_mut(),
148                    i,
149                    &sk.data.as_vec_znx(),
150                    i,
151                );
152            });
153            let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, sk_out.rank().into());
154            sk_out_prepared.prepare(module, &sk_out);
155
156            let mut atk: GLWEAutomorphismKey<Vec<u8>> = GLWEAutomorphismKey::alloc_from_infos(&atk_infos);
157            atk.decompress(module, &atk_compressed);
158
159            atk.key
160                .assert_noise(module, &sk_out_prepared, &sk.data, SIGMA);
161        }
162    }
163}