1use poulpy_hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
  4        VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace,
  5        VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
  6        VecZnxCopy, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxNormalize,
  7        VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace, VecZnxSubScalarInplace, VecZnxSwitchRing,
  8        VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, VmpPMatAlloc, VmpPrepare,
  9    },
 10    layouts::{Backend, Module, ScratchOwned},
 11    oep::{
 12        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
 13        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
 14    },
 15    source::Source,
 16};
 17
 18use crate::{
 19    encryption::SIGMA,
 20    layouts::{
 21        AutomorphismKey, AutomorphismKeyLayout, GLWEInfos, GLWESecret,
 22        compressed::{AutomorphismKeyCompressed, Decompress},
 23        prepared::{GLWESecretPrepared, PrepareAlloc},
 24    },
 25};
 26
 27pub fn test_gglwe_automorphisk_key_encrypt_sk<B>(module: &Module<B>)
 28where
 29    Module<B>: VecZnxDftAllocBytes
 30        + VecZnxBigNormalize<B>
 31        + VecZnxDftApply<B>
 32        + SvpApplyDftToDftInplace<B>
 33        + VecZnxIdftApplyConsume<B>
 34        + VecZnxNormalizeTmpBytes
 35        + VecZnxFillUniform
 36        + VecZnxSubInplace
 37        + VecZnxAddInplace
 38        + VecZnxNormalizeInplace<B>
 39        + VecZnxAddNormal
 40        + VecZnxNormalize<B>
 41        + VecZnxSub
 42        + SvpPrepare<B>
 43        + SvpPPolAllocBytes
 44        + SvpPPolAlloc<B>
 45        + VmpApplyDftToDftTmpBytes
 46        + VecZnxBigNormalizeTmpBytes
 47        + VmpApplyDftToDft<B>
 48        + VmpApplyDftToDftAdd<B>
 49        + VecZnxBigAddSmallInplace<B>
 50        + VecZnxAutomorphism
 51        + VecZnxSwitchRing
 52        + VecZnxAddScalarInplace
 53        + VecZnxAutomorphismInplace<B>
 54        + VecZnxBigAllocBytes
 55        + VecZnxBigAddInplace<B>
 56        + VecZnxSubScalarInplace
 57        + VecZnxCopy
 58        + VmpPMatAlloc<B>
 59        + VmpPrepare<B>,
 60    B: Backend
 61        + ScratchOwnedAllocImpl<B>
 62        + ScratchOwnedBorrowImpl<B>
 63        + ScratchAvailableImpl<B>
 64        + TakeScalarZnxImpl<B>
 65        + TakeVecZnxDftImpl<B>
 66        + TakeVecZnxImpl<B>
 67        + TakeSvpPPolImpl<B>
 68        + TakeVecZnxBigImpl<B>,
 69{
 70    let base2k: usize = 12;
 71    let k_ksk: usize = 60;
 72    let dsize: usize = k_ksk.div_ceil(base2k) - 1;
 73    for rank in 1_usize..3 {
 74        for di in 1..dsize + 1 {
 75            let n: usize = module.n();
 76            let dnum: usize = (k_ksk - di * base2k) / (di * base2k);
 77
 78            let atk_infos: AutomorphismKeyLayout = AutomorphismKeyLayout {
 79                n: n.into(),
 80                base2k: base2k.into(),
 81                k: k_ksk.into(),
 82                dnum: dnum.into(),
 83                dsize: di.into(),
 84                rank: rank.into(),
 85            };
 86
 87            let mut atk: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc_from_infos(&atk_infos);
 88
 89            let mut source_xs: Source = Source::new([0u8; 32]);
 90            let mut source_xe: Source = Source::new([0u8; 32]);
 91            let mut source_xa: Source = Source::new([0u8; 32]);
 92
 93            let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(AutomorphismKey::encrypt_sk_scratch_space(
 94                module, &atk_infos,
 95            ));
 96
 97            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&atk_infos);
 98            sk.fill_ternary_prob(0.5, &mut source_xs);
 99
100            let p = -5;
101
102            atk.encrypt_sk(
103                module,
104                p,
105                &sk,
106                &mut source_xa,
107                &mut source_xe,
108                scratch.borrow(),
109            );
110
111            let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
112            (0..atk.rank().into()).for_each(|i| {
113                module.vec_znx_automorphism(
114                    module.galois_element_inv(p),
115                    &mut sk_out.data.as_vec_znx_mut(),
116                    i,
117                    &sk.data.as_vec_znx(),
118                    i,
119                );
120            });
121            let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
122
123            atk.key
124                .key
125                .assert_noise(module, &sk_out_prepared, &sk.data, SIGMA);
126        }
127    }
128}
129
130pub fn test_gglwe_automorphisk_key_compressed_encrypt_sk<B>(module: &Module<B>)
131where
132    Module<B>: VecZnxDftAllocBytes
133        + VecZnxBigNormalize<B>
134        + VecZnxDftApply<B>
135        + SvpApplyDftToDftInplace<B>
136        + VecZnxIdftApplyConsume<B>
137        + VecZnxNormalizeTmpBytes
138        + VecZnxFillUniform
139        + VecZnxSubInplace
140        + VecZnxAddInplace
141        + VecZnxNormalizeInplace<B>
142        + VecZnxAddNormal
143        + VecZnxNormalize<B>
144        + VecZnxSub
145        + SvpPrepare<B>
146        + SvpPPolAllocBytes
147        + SvpPPolAlloc<B>
148        + VmpApplyDftToDftTmpBytes
149        + VecZnxBigNormalizeTmpBytes
150        + VmpApplyDftToDft<B>
151        + VmpApplyDftToDftAdd<B>
152        + VecZnxBigAddSmallInplace<B>
153        + VecZnxAutomorphism
154        + VecZnxSwitchRing
155        + VecZnxAddScalarInplace
156        + VecZnxAutomorphismInplace<B>
157        + VecZnxBigAllocBytes
158        + VecZnxBigAddInplace<B>
159        + VecZnxSubScalarInplace
160        + VecZnxCopy
161        + VmpPMatAlloc<B>
162        + VmpPrepare<B>,
163    B: Backend
164        + ScratchOwnedAllocImpl<B>
165        + ScratchOwnedBorrowImpl<B>
166        + ScratchAvailableImpl<B>
167        + TakeScalarZnxImpl<B>
168        + TakeVecZnxDftImpl<B>
169        + TakeVecZnxImpl<B>
170        + TakeSvpPPolImpl<B>
171        + TakeVecZnxBigImpl<B>,
172{
173    let base2k: usize = 12;
174    let k_ksk: usize = 60;
175    let dsize: usize = k_ksk.div_ceil(base2k) - 1;
176    for rank in 1_usize..3 {
177        for di in 1..dsize + 1 {
178            let n: usize = module.n();
179            let dnum: usize = (k_ksk - di * base2k) / (di * base2k);
180
181            let atk_infos: AutomorphismKeyLayout = AutomorphismKeyLayout {
182                n: n.into(),
183                base2k: base2k.into(),
184                k: k_ksk.into(),
185                dnum: dnum.into(),
186                dsize: di.into(),
187                rank: rank.into(),
188            };
189
190            let mut atk_compressed: AutomorphismKeyCompressed<Vec<u8>> = AutomorphismKeyCompressed::alloc_from_infos(&atk_infos);
191
192            let mut source_xs: Source = Source::new([0u8; 32]);
193            let mut source_xe: Source = Source::new([0u8; 32]);
194
195            let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(AutomorphismKey::encrypt_sk_scratch_space(
196                module, &atk_infos,
197            ));
198
199            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&atk_infos);
200            sk.fill_ternary_prob(0.5, &mut source_xs);
201
202            let p = -5;
203
204            let seed_xa: [u8; 32] = [1u8; 32];
205
206            atk_compressed.encrypt_sk(module, p, &sk, seed_xa, &mut source_xe, scratch.borrow());
207
208            let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
209            (0..atk_compressed.rank().into()).for_each(|i| {
210                module.vec_znx_automorphism(
211                    module.galois_element_inv(p),
212                    &mut sk_out.data.as_vec_znx_mut(),
213                    i,
214                    &sk.data.as_vec_znx(),
215                    i,
216                );
217            });
218            let sk_out_prepared = sk_out.prepare_alloc(module, scratch.borrow());
219
220            let mut atk: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc_from_infos(&atk_infos);
221            atk.decompress(module, &atk_compressed);
222
223            atk.key
224                .key
225                .assert_noise(module, &sk_out_prepared, &sk.data, SIGMA);
226        }
227    }
228}