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