1use backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace,
  4        VecZnxFillUniform, VecZnxSwithcDegree, VmpPMatAlloc, VmpPMatPrepare,
  5    },
  6    layouts::{Backend, Module, ScratchOwned},
  7    oep::{
  8        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
  9        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
 10    },
 11};
 12use sampling::source::Source;
 13
 14use crate::{
 15    layouts::{
 16        GGLWEAutomorphismKey, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos,
 17        prepared::{GGLWEAutomorphismKeyPrepared, GLWESecretPrepared, Prepare, PrepareAlloc},
 18    },
 19    noise::log2_std_noise_gglwe_product,
 20    trait_families::{GLWEDecryptFamily, GLWEKeyswitchFamily},
 21};
 22
 23use crate::trait_families::{GGLWEAutomorphismKeyEncryptSkFamily, GLWESecretPreparedModuleFamily};
 24
 25pub fn test_glwe_automorphism<B: Backend>(
 26    module: &Module<B>,
 27    basek: usize,
 28    p: i64,
 29    k_out: usize,
 30    k_in: usize,
 31    k_ksk: usize,
 32    digits: usize,
 33    rank: usize,
 34    sigma: f64,
 35) where
 36    Module<B>: GGLWEAutomorphismKeyEncryptSkFamily<B>
 37        + GLWESecretPreparedModuleFamily<B>
 38        + GLWEDecryptFamily<B>
 39        + GLWEKeyswitchFamily<B>
 40        + VecZnxAutomorphism
 41        + VecZnxSwithcDegree
 42        + VecZnxAddScalarInplace
 43        + VecZnxAutomorphismInplace
 44        + VmpPMatAlloc<B>
 45        + VmpPMatPrepare<B>,
 46    B: TakeVecZnxDftImpl<B>
 47        + TakeVecZnxBigImpl<B>
 48        + TakeSvpPPolImpl<B>
 49        + ScratchOwnedAllocImpl<B>
 50        + ScratchOwnedBorrowImpl<B>
 51        + ScratchAvailableImpl<B>
 52        + TakeScalarZnxImpl<B>
 53        + TakeVecZnxImpl<B>,
 54{
 55    let n: usize = module.n();
 56    let rows: usize = k_in.div_ceil(basek * digits);
 57
 58    let mut autokey: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
 59    let mut ct_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_in, rank);
 60    let mut ct_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_out, rank);
 61    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_in);
 62
 63    let mut source_xs: Source = Source::new([0u8; 32]);
 64    let mut source_xe: Source = Source::new([0u8; 32]);
 65    let mut source_xa: Source = Source::new([0u8; 32]);
 66
 67    module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_in, &mut source_xa);
 68
 69    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 70        GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, autokey.k(), rank)
 71            | GLWECiphertext::decrypt_scratch_space(module, n, basek, ct_out.k())
 72            | GLWECiphertext::encrypt_sk_scratch_space(module, n, basek, ct_in.k())
 73            | GLWECiphertext::automorphism_scratch_space(
 74                module,
 75                n,
 76                basek,
 77                ct_out.k(),
 78                ct_in.k(),
 79                autokey.k(),
 80                digits,
 81                rank,
 82            ),
 83    );
 84
 85    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
 86    sk.fill_ternary_prob(0.5, &mut source_xs);
 87    let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
 88
 89    autokey.encrypt_sk(
 90        module,
 91        p,
 92        &sk,
 93        &mut source_xa,
 94        &mut source_xe,
 95        sigma,
 96        scratch.borrow(),
 97    );
 98
 99    ct_in.encrypt_sk(
100        module,
101        &pt_want,
102        &sk_prepared,
103        &mut source_xa,
104        &mut source_xe,
105        sigma,
106        scratch.borrow(),
107    );
108
109    let mut autokey_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, B> =
110        GGLWEAutomorphismKeyPrepared::alloc(module, n, basek, k_ksk, rows, digits, rank);
111    autokey_prepared.prepare(module, &autokey, scratch.borrow());
112
113    ct_out.automorphism(module, &ct_in, &autokey_prepared, scratch.borrow());
114
115    let max_noise: f64 = log2_std_noise_gglwe_product(
116        module.n() as f64,
117        basek * digits,
118        0.5,
119        0.5,
120        0f64,
121        sigma * sigma,
122        0f64,
123        rank as f64,
124        k_in,
125        k_ksk,
126    );
127
128    module.vec_znx_automorphism_inplace(p, &mut pt_want.data, 0);
129
130    ct_out.assert_noise(module, &sk_prepared, &pt_want, max_noise + 1.0);
131}
132
133pub fn test_glwe_automorphism_inplace<B: Backend>(
134    module: &Module<B>,
135    basek: usize,
136    p: i64,
137    k_ct: usize,
138    k_ksk: usize,
139    digits: usize,
140    rank: usize,
141    sigma: f64,
142) where
143    Module<B>: GGLWEAutomorphismKeyEncryptSkFamily<B>
144        + GLWESecretPreparedModuleFamily<B>
145        + GLWEDecryptFamily<B>
146        + GLWEKeyswitchFamily<B>
147        + VecZnxAutomorphism
148        + VecZnxSwithcDegree
149        + VecZnxAddScalarInplace
150        + VecZnxAutomorphismInplace
151        + VmpPMatAlloc<B>
152        + VmpPMatPrepare<B>,
153    B: TakeVecZnxDftImpl<B>
154        + TakeVecZnxBigImpl<B>
155        + TakeSvpPPolImpl<B>
156        + ScratchOwnedAllocImpl<B>
157        + ScratchOwnedBorrowImpl<B>
158        + ScratchAvailableImpl<B>
159        + TakeScalarZnxImpl<B>
160        + TakeVecZnxImpl<B>,
161{
162    let n: usize = module.n();
163    let rows: usize = k_ct.div_ceil(basek * digits);
164
165    let mut autokey: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
166    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_ct, rank);
167    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_ct);
168
169    let mut source_xs: Source = Source::new([0u8; 32]);
170    let mut source_xe: Source = Source::new([0u8; 32]);
171    let mut source_xa: Source = Source::new([0u8; 32]);
172
173    module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_ct, &mut source_xa);
174
175    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
176        GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, autokey.k(), rank)
177            | GLWECiphertext::decrypt_scratch_space(module, n, basek, ct.k())
178            | GLWECiphertext::encrypt_sk_scratch_space(module, n, basek, ct.k())
179            | GLWECiphertext::automorphism_inplace_scratch_space(module, n, basek, ct.k(), autokey.k(), digits, 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    let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
185
186    autokey.encrypt_sk(
187        module,
188        p,
189        &sk,
190        &mut source_xa,
191        &mut source_xe,
192        sigma,
193        scratch.borrow(),
194    );
195
196    ct.encrypt_sk(
197        module,
198        &pt_want,
199        &sk_prepared,
200        &mut source_xa,
201        &mut source_xe,
202        sigma,
203        scratch.borrow(),
204    );
205
206    let mut autokey_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, B> =
207        GGLWEAutomorphismKeyPrepared::alloc(module, n, basek, k_ksk, rows, digits, rank);
208    autokey_prepared.prepare(module, &autokey, scratch.borrow());
209
210    ct.automorphism_inplace(module, &autokey_prepared, scratch.borrow());
211
212    let max_noise: f64 = log2_std_noise_gglwe_product(
213        module.n() as f64,
214        basek * digits,
215        0.5,
216        0.5,
217        0f64,
218        sigma * sigma,
219        0f64,
220        rank as f64,
221        k_ct,
222        k_ksk,
223    );
224
225    module.vec_znx_automorphism_inplace(p, &mut pt_want.data, 0);
226
227    ct.assert_noise(module, &sk_prepared, &pt_want, max_noise + 1.0);
228}