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    GGLWEEncryptSk, GGLWEKeyswitch, GLWESwitchingKeyCompressedEncryptSk, GLWESwitchingKeyEncryptSk, ScratchTakeCore,
  9    encryption::SIGMA,
 10    layouts::{
 11        AutomorphismKey, AutomorphismKeyDecompress, AutomorphismKeyLayout, GLWEInfos, GLWESecret, GLWESecretPrepare,
 12        GLWESecretPreparedAlloc, GLWESwitchingKeyDecompress, compressed::AutomorphismKeyCompressed, prepared::GLWESecretPrepared,
 13    },
 14    noise::GGLWENoise,
 15};
 16
 17pub fn test_gglwe_automorphism_key_encrypt_sk<BE: Backend>(module: &Module<BE>)
 18where
 19    Module<BE>: GGLWEEncryptSk<BE>
 20        + GGLWEKeyswitch<BE>
 21        + GLWESecretPreparedAlloc<BE>
 22        + GLWESecretPrepare<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                .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>: GGLWEEncryptSk<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(AutomorphismKey::encrypt_sk_tmp_bytes(module, &atk_infos));
129
130            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&atk_infos);
131            sk.fill_ternary_prob(0.5, &mut source_xs);
132
133            let p = -5;
134
135            let seed_xa: [u8; 32] = [1u8; 32];
136
137            atk_compressed.encrypt_sk(module, p, &sk, seed_xa, &mut source_xe, scratch.borrow());
138
139            let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
140            (0..atk_compressed.rank().into()).for_each(|i| {
141                module.vec_znx_automorphism(
142                    module.galois_element_inv(p),
143                    &mut sk_out.data.as_vec_znx_mut(),
144                    i,
145                    &sk.data.as_vec_znx(),
146                    i,
147                );
148            });
149            let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, sk_out.rank().into());
150            sk_out_prepared.prepare(module, &sk_out);
151
152            let mut atk: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc_from_infos(&atk_infos);
153            atk.decompress(module, &atk_compressed);
154
155            atk.key
156                .key
157                .assert_noise(module, &sk_out_prepared, &sk.data, SIGMA);
158        }
159    }
160}