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