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