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