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