1use backend::{
  2    hal::{
  3        api::{
  4            MatZnxAlloc, ModuleNew, ScalarZnxAlloc, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddNormal,
  5            VecZnxAddScalarInplace, VecZnxAlloc, VecZnxAllocBytes, VecZnxEncodeCoeffsi64, VecZnxFillUniform, VecZnxRotateInplace,
  6            VecZnxSub, VecZnxSwithcDegree, ZnxView,
  7        },
  8        layouts::{Backend, Module, ScratchOwned},
  9        oep::{
 10            ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeVecZnxBigImpl, TakeVecZnxDftImpl,
 11            TakeVecZnxDftSliceImpl, TakeVecZnxImpl, TakeVecZnxSliceImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
 12        },
 13    },
 14    implementation::cpu_spqlios::FFT64,
 15};
 16use sampling::source::Source;
 17
 18use crate::{
 19    BlindRotationKeyCGGIExecLayoutFamily, CCGIBlindRotationFamily, GLWECiphertext, GLWEDecryptFamily, GLWEPlaintext, GLWESecret,
 20    GLWESecretExec, GLWESecretFamily, Infos, LWECiphertext, LWESecret,
 21    blind_rotation::{
 22        cggi::{cggi_blind_rotate, cggi_blind_rotate_scratch_space, negate_and_mod_switch_2n},
 23        key::{BlindRotationKeyCGGI, BlindRotationKeyCGGIExec},
 24        lut::LookUpTable,
 25    },
 26    lwe::{LWEPlaintext, ciphertext::LWECiphertextToRef},
 27};
 28
 29#[test]
 30fn standard() {
 31    let module: Module<FFT64> = Module::<FFT64>::new(512);
 32    blind_rotatio_test(&module, 224, 1, 1);
 33}
 34
 35#[test]
 36fn block_binary() {
 37    let module: Module<FFT64> = Module::<FFT64>::new(512);
 38    blind_rotatio_test(&module, 224, 7, 1);
 39}
 40
 41#[test]
 42fn block_binary_extended() {
 43    let module: Module<FFT64> = Module::<FFT64>::new(512);
 44    blind_rotatio_test(&module, 224, 7, 2);
 45}
 46
 47pub(crate) trait CGGITestModuleFamily<B: Backend> = CCGIBlindRotationFamily<B>
 48    + GLWESecretFamily<B>
 49    + GLWEDecryptFamily<B>
 50    + BlindRotationKeyCGGIExecLayoutFamily<B>
 51    + VecZnxAlloc
 52    + ScalarZnxAlloc
 53    + VecZnxFillUniform
 54    + VecZnxAddNormal
 55    + VecZnxAllocBytes
 56    + VecZnxAddScalarInplace
 57    + VecZnxEncodeCoeffsi64
 58    + VecZnxRotateInplace
 59    + VecZnxSwithcDegree
 60    + MatZnxAlloc
 61    + VecZnxSub;
 62pub(crate) trait CGGITestScratchFamily<B: Backend> = VecZnxDftAllocBytesImpl<B>
 63    + VecZnxBigAllocBytesImpl<B>
 64    + ScratchOwnedAllocImpl<B>
 65    + ScratchOwnedBorrowImpl<B>
 66    + TakeVecZnxDftImpl<B>
 67    + TakeVecZnxBigImpl<B>
 68    + TakeVecZnxDftSliceImpl<B>
 69    + ScratchAvailableImpl<B>
 70    + TakeVecZnxImpl<B>
 71    + TakeVecZnxSliceImpl<B>;
 72
 73fn blind_rotatio_test<B: Backend>(module: &Module<B>, n_lwe: usize, block_size: usize, extension_factor: usize)
 74where
 75    Module<B>: CGGITestModuleFamily<B>,
 76    B: CGGITestScratchFamily<B>,
 77{
 78    let basek: usize = 19;
 79
 80    let k_lwe: usize = 24;
 81    let k_brk: usize = 3 * basek;
 82    let rows_brk: usize = 2; // Ensures first limb is noise-free.
 83    let k_lut: usize = 1 * basek;
 84    let k_res: usize = 2 * basek;
 85    let rank: usize = 1;
 86
 87    let message_modulus: usize = 1 << 4;
 88
 89    let mut source_xs: Source = Source::new([2u8; 32]);
 90    let mut source_xe: Source = Source::new([2u8; 32]);
 91    let mut source_xa: Source = Source::new([1u8; 32]);
 92
 93    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank);
 94    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
 95    let sk_glwe_dft: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_glwe);
 96
 97    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
 98    sk_lwe.fill_binary_block(block_size, &mut source_xs);
 99
100    let mut scratch: ScratchOwned<B> = ScratchOwned::<B>::alloc(BlindRotationKeyCGGI::generate_from_sk_scratch_space(
101        module, basek, k_brk, rank,
102    ));
103
104    let mut scratch_br: ScratchOwned<B> = ScratchOwned::<B>::alloc(cggi_blind_rotate_scratch_space(
105        module,
106        block_size,
107        extension_factor,
108        basek,
109        k_res,
110        k_brk,
111        rows_brk,
112        rank,
113    ));
114
115    let mut brk: BlindRotationKeyCGGI<Vec<u8>> = BlindRotationKeyCGGI::alloc(module, n_lwe, basek, k_brk, rows_brk, rank);
116
117    brk.generate_from_sk(
118        module,
119        &sk_glwe_dft,
120        &sk_lwe,
121        &mut source_xa,
122        &mut source_xe,
123        3.2,
124        scratch.borrow(),
125    );
126
127    let mut lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe);
128
129    let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe);
130
131    let x: i64 = 2;
132    let bits: usize = 8;
133
134    module.encode_coeff_i64(basek, &mut pt_lwe.data, 0, bits, 0, x, bits);
135
136    lwe.encrypt_sk(
137        module,
138        &pt_lwe,
139        &sk_lwe,
140        &mut source_xa,
141        &mut source_xe,
142        3.2,
143    );
144
145    let mut f: Vec<i64> = vec![0i64; message_modulus];
146    f.iter_mut()
147        .enumerate()
148        .for_each(|(i, x)| *x = 2 * (i as i64) + 1);
149
150    let mut lut: LookUpTable = LookUpTable::alloc(module, basek, k_lut, extension_factor);
151    lut.set(module, &f, message_modulus);
152
153    let mut res: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_res, rank);
154
155    let brk_exec: BlindRotationKeyCGGIExec<Vec<u8>, B> = BlindRotationKeyCGGIExec::from(module, &brk, scratch_br.borrow());
156
157    cggi_blind_rotate(module, &mut res, &lwe, &lut, &brk_exec, scratch_br.borrow());
158
159    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_res);
160
161    res.decrypt(module, &mut pt_have, &sk_glwe_dft, scratch.borrow());
162
163    let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
164
165    negate_and_mod_switch_2n(2 * lut.domain_size(), &mut lwe_2n, &lwe.to_ref());
166
167    let pt_want: i64 = (lwe_2n[0]
168        + lwe_2n[1..]
169            .iter()
170            .zip(sk_lwe.data.at(0, 0))
171            .map(|(x, y)| x * y)
172            .sum::<i64>())
173        & (2 * lut.domain_size() - 1) as i64;
174
175    lut.rotate(module, pt_want);
176
177    // First limb should be exactly equal (test are parameterized such that the noise does not reach
178    // the first limb)
179    assert_eq!(pt_have.data.at(0, 0), lut.data[0].at(0, 0));
180}