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