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