1use poulpy_hal::{
  2    api::{ScratchOwnedAlloc, ScratchOwnedBorrow},
  3    layouts::{Backend, Scratch, ScratchOwned, ZnxView},
  4    source::Source,
  5};
  6
  7use crate::tfhe::blind_rotation::{
  8    BlindRotationAlgo, BlindRotationExecute, BlindRotationKey, BlindRotationKeyEncryptSk, BlindRotationKeyFactory,
  9    BlindRotationKeyLayout, BlindRotationKeyPrepared, BlindRotationKeyPreparedFactory, LookUpTableLayout, LookupTable,
 10    LookupTableFactory, mod_switch_2n,
 11};
 12
 13use poulpy_core::{
 14    GLWEDecrypt, LWEEncryptSk, ScratchTakeCore,
 15    layouts::{
 16        GLWE, GLWELayout, GLWEPlaintext, GLWESecret, GLWESecretPreparedFactory, LWE, LWEInfos, LWELayout, LWEPlaintext,
 17        LWESecret, LWEToRef, prepared::GLWESecretPrepared,
 18    },
 19};
 20
 21pub fn test_blind_rotation<BRA: BlindRotationAlgo, M, BE: Backend>(
 22    module: &M,
 23    n_lwe: usize,
 24    block_size: usize,
 25    extension_factor: usize,
 26) where
 27    M: BlindRotationKeyEncryptSk<BRA, BE>
 28        + BlindRotationKeyPreparedFactory<BRA, BE>
 29        + BlindRotationExecute<BRA, BE>
 30        + LookupTableFactory
 31        + GLWESecretPreparedFactory<BE>
 32        + GLWEDecrypt<BE>
 33        + LWEEncryptSk<BE>,
 34    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
 35    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 36    Scratch<BE>: ScratchTakeCore<BE>,
 37{
 38    let n_glwe: usize = module.n();
 39    let base2k: usize = 19;
 40    let k_lwe: usize = 24;
 41    let k_brk: usize = 3 * base2k;
 42    let rows_brk: usize = 2; // Ensures first limb is noise-free.
 43    let k_lut: usize = base2k;
 44    let k_res: usize = 2 * base2k;
 45    let rank: usize = 1;
 46
 47    let log_message_modulus: usize = 4;
 48
 49    let message_modulus: usize = 1 << log_message_modulus;
 50
 51    let mut source_xs: Source = Source::new([2u8; 32]);
 52    let mut source_xe: Source = Source::new([2u8; 32]);
 53    let mut source_xa: Source = Source::new([1u8; 32]);
 54
 55    let brk_infos: BlindRotationKeyLayout = BlindRotationKeyLayout {
 56        n_glwe: n_glwe.into(),
 57        n_lwe: n_lwe.into(),
 58        base2k: base2k.into(),
 59        k: k_brk.into(),
 60        dnum: rows_brk.into(),
 61        rank: rank.into(),
 62    };
 63
 64    let glwe_infos: GLWELayout = GLWELayout {
 65        n: n_glwe.into(),
 66        base2k: base2k.into(),
 67        k: k_res.into(),
 68        rank: rank.into(),
 69    };
 70
 71    let lwe_infos: LWELayout = LWELayout {
 72        n: n_lwe.into(),
 73        k: k_lwe.into(),
 74        base2k: base2k.into(),
 75    };
 76
 77    let mut scratch: ScratchOwned<BE> = ScratchOwned::<BE>::alloc(BlindRotationKey::encrypt_sk_tmp_bytes(module, &brk_infos));
 78
 79    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_infos);
 80    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
 81    let mut sk_glwe_dft: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc_from_infos(module, &glwe_infos);
 82    sk_glwe_dft.prepare(module, &sk_glwe);
 83
 84    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe.into());
 85    sk_lwe.fill_binary_block(block_size, &mut source_xs);
 86
 87    let mut scratch_br: ScratchOwned<BE> = ScratchOwned::<BE>::alloc(BlindRotationKeyPrepared::execute_tmp_bytes(
 88        module,
 89        block_size,
 90        extension_factor,
 91        &glwe_infos,
 92        &brk_infos,
 93    ));
 94
 95    let mut brk: BlindRotationKey<Vec<u8>, BRA> = BlindRotationKey::<Vec<u8>, BRA>::alloc(&brk_infos);
 96
 97    brk.encrypt_sk(
 98        module,
 99        &sk_glwe_dft,
100        &sk_lwe,
101        &mut source_xa,
102        &mut source_xe,
103        scratch.borrow(),
104    );
105
106    let mut lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_infos);
107
108    let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc_from_infos(&lwe_infos);
109
110    let x: i64 = 15 % (message_modulus as i64);
111
112    pt_lwe.encode_i64(x, (log_message_modulus + 1).into());
113
114    lwe.encrypt_sk(
115        module,
116        &pt_lwe,
117        &sk_lwe,
118        &mut source_xa,
119        &mut source_xe,
120        scratch.borrow(),
121    );
122
123    let f = |x: i64| -> i64 { 2 * x + 1 };
124
125    let mut f_vec: Vec<i64> = vec![0i64; message_modulus];
126    f_vec
127        .iter_mut()
128        .enumerate()
129        .for_each(|(i, x)| *x = f(i as i64));
130
131    let lut_infos = LookUpTableLayout {
132        n: module.n().into(),
133        extension_factor,
134        k: k_lut.into(),
135        base2k: base2k.into(),
136    };
137
138    let mut lut: LookupTable = LookupTable::alloc(&lut_infos);
139    lut.set(module, &f_vec, log_message_modulus + 1);
140
141    let mut res: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_infos);
142
143    let mut brk_prepared: BlindRotationKeyPrepared<Vec<u8>, BRA, BE> = BlindRotationKeyPrepared::alloc(module, &brk);
144    brk_prepared.prepare(module, &brk, scratch_br.borrow());
145
146    brk_prepared.execute(module, &mut res, &lwe, &lut, scratch_br.borrow());
147
148    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_infos);
149
150    res.decrypt(module, &mut pt_have, &sk_glwe_dft, scratch.borrow());
151
152    let mut lwe_2n: Vec<i64> = vec![0i64; (lwe.n() + 1).into()]; // TODO: from scratch space
153
154    mod_switch_2n(
155        2 * lut.domain_size(),
156        &mut lwe_2n,
157        &lwe.to_ref(),
158        lut.rotation_direction(),
159    );
160
161    let pt_want: i64 = (lwe_2n[0]
162        + lwe_2n[1..]
163            .iter()
164            .zip(sk_lwe.raw())
165            .map(|(x, y)| x * y)
166            .sum::<i64>())
167        & (2 * lut.domain_size() - 1) as i64;
168
169    lut.rotate(module, pt_want);
170
171    // First limb should be exactly equal (test are parameterized such that the noise does not reach
172    // the first limb)
173    assert_eq!(pt_have.data.at(0, 0), lut.data[0].at(0, 0));
174
175    // Verify that it effectively compute f(x)
176    let mut have: i64 = pt_have.decode_coeff_i64((log_message_modulus + 1).into(), 0);
177
178    // Get positive representative and assert equality
179    have = (have + message_modulus as i64) % (message_modulus as i64);
180
181    assert_eq!(have, f(x) % (message_modulus as i64));
182}