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