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, VecZnxSubABInplace, 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        GGLWEAutomorphismKey, GLWESecret,
 22        compressed::{Decompress, GGLWEAutomorphismKeyCompressed},
 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        + VecZnxSubABInplace
 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 basek: usize = 12;
 71    let k_ksk: usize = 60;
 72    let digits: usize = k_ksk.div_ceil(basek) - 1;
 73    (1..3).for_each(|rank| {
 74        (1..digits + 1).for_each(|di| {
 75            let n: usize = module.n();
 76            let rows: usize = (k_ksk - di * basek) / (di * basek);
 77
 78            let mut atk: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, di, rank);
 79
 80            let mut source_xs: Source = Source::new([0u8; 32]);
 81            let mut source_xe: Source = Source::new([0u8; 32]);
 82            let mut source_xa: Source = Source::new([0u8; 32]);
 83
 84            let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWEAutomorphismKey::encrypt_sk_scratch_space(
 85                module, basek, k_ksk, rank,
 86            ));
 87
 88            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
 89            sk.fill_ternary_prob(0.5, &mut source_xs);
 90
 91            let p = -5;
 92
 93            atk.encrypt_sk(
 94                module,
 95                p,
 96                &sk,
 97                &mut source_xa,
 98                &mut source_xe,
 99                scratch.borrow(),
100            );
101
102            let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
103            (0..atk.rank()).for_each(|i| {
104                module.vec_znx_automorphism(
105                    module.galois_element_inv(p),
106                    &mut sk_out.data.as_vec_znx_mut(),
107                    i,
108                    &sk.data.as_vec_znx(),
109                    i,
110                );
111            });
112            let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
113
114            atk.key
115                .key
116                .assert_noise(module, &sk_out_prepared, &sk.data, SIGMA);
117        });
118    });
119}
120
121pub fn test_gglwe_automorphisk_key_compressed_encrypt_sk<B>(module: &Module<B>)
122where
123    Module<B>: VecZnxDftAllocBytes
124        + VecZnxBigNormalize<B>
125        + VecZnxDftApply<B>
126        + SvpApplyDftToDftInplace<B>
127        + VecZnxIdftApplyConsume<B>
128        + VecZnxNormalizeTmpBytes
129        + VecZnxFillUniform
130        + VecZnxSubABInplace
131        + VecZnxAddInplace
132        + VecZnxNormalizeInplace<B>
133        + VecZnxAddNormal
134        + VecZnxNormalize<B>
135        + VecZnxSub
136        + SvpPrepare<B>
137        + SvpPPolAllocBytes
138        + SvpPPolAlloc<B>
139        + VmpApplyDftToDftTmpBytes
140        + VecZnxBigNormalizeTmpBytes
141        + VmpApplyDftToDft<B>
142        + VmpApplyDftToDftAdd<B>
143        + VecZnxBigAddSmallInplace<B>
144        + VecZnxAutomorphism
145        + VecZnxSwitchRing
146        + VecZnxAddScalarInplace
147        + VecZnxAutomorphismInplace<B>
148        + VecZnxBigAllocBytes
149        + VecZnxBigAddInplace<B>
150        + VecZnxSubScalarInplace
151        + VecZnxCopy
152        + VmpPMatAlloc<B>
153        + VmpPrepare<B>,
154    B: Backend
155        + ScratchOwnedAllocImpl<B>
156        + ScratchOwnedBorrowImpl<B>
157        + ScratchAvailableImpl<B>
158        + TakeScalarZnxImpl<B>
159        + TakeVecZnxDftImpl<B>
160        + TakeVecZnxImpl<B>
161        + TakeSvpPPolImpl<B>
162        + TakeVecZnxBigImpl<B>,
163{
164    let basek: usize = 12;
165    let k_ksk: usize = 60;
166    let digits: usize = k_ksk.div_ceil(basek) - 1;
167    (1..3).for_each(|rank| {
168        (1..digits + 1).for_each(|di| {
169            let n: usize = module.n();
170            let rows: usize = (k_ksk - di * basek) / (di * basek);
171
172            let mut atk_compressed: GGLWEAutomorphismKeyCompressed<Vec<u8>> =
173                GGLWEAutomorphismKeyCompressed::alloc(n, basek, k_ksk, rows, di, rank);
174
175            let mut source_xs: Source = Source::new([0u8; 32]);
176            let mut source_xe: Source = Source::new([0u8; 32]);
177
178            let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWEAutomorphismKey::encrypt_sk_scratch_space(
179                module, basek, k_ksk, rank,
180            ));
181
182            let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
183            sk.fill_ternary_prob(0.5, &mut source_xs);
184
185            let p = -5;
186
187            let seed_xa: [u8; 32] = [1u8; 32];
188
189            atk_compressed.encrypt_sk(module, p, &sk, seed_xa, &mut source_xe, scratch.borrow());
190
191            let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
192            (0..atk_compressed.rank()).for_each(|i| {
193                module.vec_znx_automorphism(
194                    module.galois_element_inv(p),
195                    &mut sk_out.data.as_vec_znx_mut(),
196                    i,
197                    &sk.data.as_vec_znx(),
198                    i,
199                );
200            });
201            let sk_out_prepared = sk_out.prepare_alloc(module, scratch.borrow());
202
203            let mut atk: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, di, rank);
204            atk.decompress(module, &atk_compressed);
205
206            atk.key
207                .key
208                .assert_noise(module, &sk_out_prepared, &sk.data, SIGMA);
209        });
210    });
211}