1use backend::hal::{
  2    api::{
  3        MatZnxAlloc, ScalarZnxAlloc, ScalarZnxAllocBytes, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace,
  4        VecZnxAlloc, VecZnxAllocBytes, VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxFillUniform, VecZnxStd,
  5        VecZnxSwithcDegree,
  6    },
  7    layouts::{Backend, Module, ScratchOwned},
  8    oep::{
  9        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
 10        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
 11    },
 12};
 13use sampling::source::Source;
 14
 15use crate::{
 16    AutomorphismKey, AutomorphismKeyEncryptSkFamily, AutomorphismKeyExec, GGLWEExecLayoutFamily, GLWECiphertext,
 17    GLWEDecryptFamily, GLWEKeyswitchFamily, GLWEPlaintext, GLWESecret, GLWESecretExec, Infos,
 18    noise::log2_std_noise_gglwe_product,
 19};
 20
 21pub(crate) trait AutomorphismTestModuleFamily<B: Backend> = AutomorphismKeyEncryptSkFamily<B>
 22    + GLWEDecryptFamily<B>
 23    + GGLWEExecLayoutFamily<B>
 24    + GLWEKeyswitchFamily<B>
 25    + MatZnxAlloc
 26    + VecZnxAlloc
 27    + ScalarZnxAllocBytes
 28    + VecZnxAllocBytes
 29    + VecZnxAutomorphism
 30    + VecZnxSwithcDegree
 31    + ScalarZnxAlloc
 32    + VecZnxAddScalarInplace
 33    + VecZnxAutomorphismInplace
 34    + VecZnxStd;
 35pub(crate) trait AutomorphismTestScratchFamily<B: Backend> = TakeVecZnxDftImpl<B>
 36    + TakeVecZnxBigImpl<B>
 37    + TakeSvpPPolImpl<B>
 38    + ScratchOwnedAllocImpl<B>
 39    + ScratchOwnedBorrowImpl<B>
 40    + ScratchAvailableImpl<B>
 41    + TakeScalarZnxImpl<B>
 42    + TakeVecZnxImpl<B>;
 43
 44pub(crate) fn test_automorphism<B: Backend>(
 45    module: &Module<B>,
 46    basek: usize,
 47    p: i64,
 48    k_out: usize,
 49    k_in: usize,
 50    k_ksk: usize,
 51    digits: usize,
 52    rank: usize,
 53    sigma: f64,
 54) where
 55    Module<B>: AutomorphismTestModuleFamily<B>,
 56    B: AutomorphismTestScratchFamily<B>,
 57{
 58    let rows: usize = k_in.div_ceil(basek * digits);
 59
 60    let mut autokey: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(module, basek, k_ksk, rows, digits, rank);
 61    let mut ct_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_in, rank);
 62    let mut ct_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_out, rank);
 63    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_in);
 64
 65    let mut source_xs: Source = Source::new([0u8; 32]);
 66    let mut source_xe: Source = Source::new([0u8; 32]);
 67    let mut source_xa: Source = Source::new([0u8; 32]);
 68
 69    module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_in, &mut source_xa);
 70
 71    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 72        AutomorphismKey::encrypt_sk_scratch_space(module, basek, autokey.k(), rank)
 73            | GLWECiphertext::decrypt_scratch_space(module, basek, ct_out.k())
 74            | GLWECiphertext::encrypt_sk_scratch_space(module, basek, ct_in.k())
 75            | GLWECiphertext::automorphism_scratch_space(
 76                module,
 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(module, 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: AutomorphismKeyExec<Vec<u8>, B> = AutomorphismKeyExec::alloc(module, basek, k_ksk, rows, digits, rank);
111    autokey_exec.prepare(module, &autokey, scratch.borrow());
112
113    ct_out.automorphism(module, &ct_in, &autokey_exec, 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_exec, &pt_want, max_noise + 1.0);
131}
132
133pub(crate) fn test_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>: AutomorphismTestModuleFamily<B>,
144    B: AutomorphismTestScratchFamily<B>,
145{
146    let rows: usize = k_ct.div_ceil(basek * digits);
147
148    let mut autokey: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(module, basek, k_ksk, rows, digits, rank);
149    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_ct, rank);
150    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_ct);
151
152    let mut source_xs: Source = Source::new([0u8; 32]);
153    let mut source_xe: Source = Source::new([0u8; 32]);
154    let mut source_xa: Source = Source::new([0u8; 32]);
155
156    module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_ct, &mut source_xa);
157
158    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
159        AutomorphismKey::encrypt_sk_scratch_space(module, basek, autokey.k(), rank)
160            | GLWECiphertext::decrypt_scratch_space(module, basek, ct.k())
161            | GLWECiphertext::encrypt_sk_scratch_space(module, basek, ct.k())
162            | GLWECiphertext::automorphism_inplace_scratch_space(module, basek, ct.k(), autokey.k(), digits, rank),
163    );
164
165    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank);
166    sk.fill_ternary_prob(0.5, &mut source_xs);
167    let sk_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk);
168
169    autokey.encrypt_sk(
170        module,
171        p,
172        &sk,
173        &mut source_xa,
174        &mut source_xe,
175        sigma,
176        scratch.borrow(),
177    );
178
179    ct.encrypt_sk(
180        module,
181        &pt_want,
182        &sk_exec,
183        &mut source_xa,
184        &mut source_xe,
185        sigma,
186        scratch.borrow(),
187    );
188
189    let mut autokey_exec: AutomorphismKeyExec<Vec<u8>, B> = AutomorphismKeyExec::alloc(module, basek, k_ksk, rows, digits, rank);
190    autokey_exec.prepare(module, &autokey, scratch.borrow());
191
192    ct.automorphism_inplace(module, &autokey_exec, scratch.borrow());
193
194    let max_noise: f64 = log2_std_noise_gglwe_product(
195        module.n() as f64,
196        basek * digits,
197        0.5,
198        0.5,
199        0f64,
200        sigma * sigma,
201        0f64,
202        rank as f64,
203        k_ct,
204        k_ksk,
205    );
206
207    module.vec_znx_automorphism_inplace(p, &mut pt_want.data, 0);
208
209    ct.assert_noise(module, &sk_exec, &pt_want, max_noise + 1.0);
210}