1use backend::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};
 16use sampling::source::Source;
 17
 18use crate::layouts::{
 19    GGLWEAutomorphismKey, GLWESecret,
 20    compressed::{Decompress, GGLWEAutomorphismKeyCompressed},
 21    prepared::{GLWESecretPrepared, PrepareAlloc},
 22};
 23
 24pub fn test_gglwe_automorphisk_key_encrypt_sk<B: Backend>(
 25    module: &Module<B>,
 26    basek: usize,
 27    k_ksk: usize,
 28    digits: usize,
 29    rank: usize,
 30    sigma: f64,
 31) where
 32    Module<B>: VecZnxDftAllocBytes
 33        + VecZnxBigNormalize<B>
 34        + VecZnxDftFromVecZnx<B>
 35        + SvpApplyInplace<B>
 36        + VecZnxDftToVecZnxBigConsume<B>
 37        + VecZnxNormalizeTmpBytes
 38        + VecZnxFillUniform
 39        + VecZnxSubABInplace
 40        + VecZnxAddInplace
 41        + VecZnxNormalizeInplace<B>
 42        + VecZnxAddNormal
 43        + VecZnxNormalize<B>
 44        + VecZnxSub
 45        + SvpPrepare<B>
 46        + SvpPPolAllocBytes
 47        + SvpPPolAlloc<B>
 48        + VmpApplyTmpBytes
 49        + VecZnxBigNormalizeTmpBytes
 50        + VmpApply<B>
 51        + VmpApplyAdd<B>
 52        + VecZnxBigAddSmallInplace<B>
 53        + VecZnxAutomorphism
 54        + VecZnxSwithcDegree
 55        + VecZnxAddScalarInplace
 56        + VecZnxAutomorphismInplace
 57        + VecZnxBigAllocBytes
 58        + VecZnxBigAddInplace<B>
 59        + VecZnxSubScalarInplace
 60        + VecZnxCopy
 61        + VmpPMatAlloc<B>
 62        + VmpPrepare<B>,
 63    B: ScratchOwnedAllocImpl<B>
 64        + ScratchOwnedBorrowImpl<B>
 65        + ScratchAvailableImpl<B>
 66        + TakeScalarZnxImpl<B>
 67        + TakeVecZnxDftImpl<B>
 68        + TakeVecZnxImpl<B>
 69        + TakeSvpPPolImpl<B>
 70        + TakeVecZnxBigImpl<B>,
 71{
 72    let n: usize = module.n();
 73    let rows: usize = (k_ksk - digits * basek) / (digits * basek);
 74
 75    let mut atk: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
 76
 77    let mut source_xs: Source = Source::new([0u8; 32]);
 78    let mut source_xe: Source = Source::new([0u8; 32]);
 79    let mut source_xa: Source = Source::new([0u8; 32]);
 80
 81    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWEAutomorphismKey::encrypt_sk_scratch_space(
 82        module, n, basek, k_ksk, rank,
 83    ));
 84
 85    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
 86    sk.fill_ternary_prob(0.5, &mut source_xs);
 87
 88    let p = -5;
 89
 90    atk.encrypt_sk(
 91        module,
 92        p,
 93        &sk,
 94        &mut source_xa,
 95        &mut source_xe,
 96        sigma,
 97        scratch.borrow(),
 98    );
 99
100    let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
101    (0..atk.rank()).for_each(|i| {
102        module.vec_znx_automorphism(
103            module.galois_element_inv(p),
104            &mut sk_out.data.as_vec_znx_mut(),
105            i,
106            &sk.data.as_vec_znx(),
107            i,
108        );
109    });
110    let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
111
112    atk.key
113        .key
114        .assert_noise(module, &sk_out_prepared, &sk.data, sigma);
115}
116
117pub fn test_gglwe_automorphisk_key_compressed_encrypt_sk<B: Backend>(
118    module: &Module<B>,
119    basek: usize,
120    k_ksk: usize,
121    digits: usize,
122    rank: usize,
123    sigma: f64,
124) where
125    Module<B>: VecZnxDftAllocBytes
126        + VecZnxBigNormalize<B>
127        + VecZnxDftFromVecZnx<B>
128        + SvpApplyInplace<B>
129        + VecZnxDftToVecZnxBigConsume<B>
130        + VecZnxNormalizeTmpBytes
131        + VecZnxFillUniform
132        + VecZnxSubABInplace
133        + VecZnxAddInplace
134        + VecZnxNormalizeInplace<B>
135        + VecZnxAddNormal
136        + VecZnxNormalize<B>
137        + VecZnxSub
138        + SvpPrepare<B>
139        + SvpPPolAllocBytes
140        + SvpPPolAlloc<B>
141        + VmpApplyTmpBytes
142        + VecZnxBigNormalizeTmpBytes
143        + VmpApply<B>
144        + VmpApplyAdd<B>
145        + VecZnxBigAddSmallInplace<B>
146        + VecZnxAutomorphism
147        + VecZnxSwithcDegree
148        + VecZnxAddScalarInplace
149        + VecZnxAutomorphismInplace
150        + VecZnxBigAllocBytes
151        + VecZnxBigAddInplace<B>
152        + VecZnxSubScalarInplace
153        + VecZnxCopy
154        + VmpPMatAlloc<B>
155        + VmpPrepare<B>,
156    B: ScratchOwnedAllocImpl<B>
157        + ScratchOwnedBorrowImpl<B>
158        + ScratchAvailableImpl<B>
159        + TakeScalarZnxImpl<B>
160        + TakeVecZnxDftImpl<B>
161        + TakeVecZnxImpl<B>
162        + TakeSvpPPolImpl<B>
163        + TakeVecZnxBigImpl<B>,
164{
165    let n: usize = module.n();
166    let rows: usize = (k_ksk - digits * basek) / (digits * basek);
167
168    let mut atk_compressed: GGLWEAutomorphismKeyCompressed<Vec<u8>> =
169        GGLWEAutomorphismKeyCompressed::alloc(n, basek, k_ksk, rows, digits, rank);
170
171    let mut source_xs: Source = Source::new([0u8; 32]);
172    let mut source_xe: Source = Source::new([0u8; 32]);
173
174    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWEAutomorphismKey::encrypt_sk_scratch_space(
175        module, n, basek, k_ksk, rank,
176    ));
177
178    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
179    sk.fill_ternary_prob(0.5, &mut source_xs);
180
181    let p = -5;
182
183    let seed_xa: [u8; 32] = [1u8; 32];
184
185    atk_compressed.encrypt_sk(
186        module,
187        p,
188        &sk,
189        seed_xa,
190        &mut source_xe,
191        sigma,
192        scratch.borrow(),
193    );
194
195    let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
196    (0..atk_compressed.rank()).for_each(|i| {
197        module.vec_znx_automorphism(
198            module.galois_element_inv(p),
199            &mut sk_out.data.as_vec_znx_mut(),
200            i,
201            &sk.data.as_vec_znx(),
202            i,
203        );
204    });
205    let sk_out_prepared = sk_out.prepare_alloc(module, scratch.borrow());
206
207    let mut atk: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
208    atk.decompress(module, &atk_compressed);
209
210    atk.key
211        .key
212        .assert_noise(module, &sk_out_prepared, &sk.data, sigma);
213}