1use std::collections::HashMap;
  2
  3use poulpy_core::{
  4    GGSWEncryptSk, GLWEDecrypt, GLWEEncryptSk, ScratchTakeCore,
  5    layouts::{
  6        Base2K, Dnum, Dsize, GGSWLayout, GGSWPreparedFactory, GLWE, GLWELayout, GLWEPlaintext, GLWESecretPrepared,
  7        GLWESecretPreparedFactory, Rank, TorusPrecision,
  8    },
  9};
 10use poulpy_hal::{
 11    api::{ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow},
 12    layouts::{Backend, Module, Scratch, ScratchOwned},
 13    source::Source,
 14};
 15use rand::RngCore;
 16
 17use crate::tfhe::{
 18    bdd_arithmetic::{
 19        FheUintPrepared, GLWEBlinSelection,
 20        tests::test_suite::{TEST_BASE2K, TEST_RANK, TestContext},
 21    },
 22    blind_rotation::BlindRotationAlgo,
 23};
 24
 25pub fn test_glwe_blind_selection<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
 26where
 27    Module<BE>: ModuleNew<BE>
 28        + GLWESecretPreparedFactory<BE>
 29        + GGSWPreparedFactory<BE>
 30        + GGSWEncryptSk<BE>
 31        + GLWEBlinSelection<u32, BE>
 32        + GLWEDecrypt<BE>
 33        + GLWEEncryptSk<BE>,
 34    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 35    Scratch<BE>: ScratchTakeCore<BE>,
 36{
 37    let module: &Module<BE> = &test_context.module;
 38    let sk_glwe_prep: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
 39
 40    let base2k: Base2K = TEST_BASE2K.into();
 41    let rank: Rank = TEST_RANK.into();
 42    let k_glwe: TorusPrecision = TorusPrecision(26);
 43    let k_ggsw: TorusPrecision = TorusPrecision(39);
 44    let dnum: Dnum = Dnum(3);
 45
 46    let glwe_infos: GLWELayout = GLWELayout {
 47        n: module.n().into(),
 48        base2k,
 49        k: k_glwe,
 50        rank,
 51    };
 52    let ggsw_infos: GGSWLayout = GGSWLayout {
 53        n: module.n().into(),
 54        base2k,
 55        k: k_ggsw,
 56        rank,
 57        dnum,
 58        dsize: Dsize(1),
 59    };
 60
 61    let mut source: Source = Source::new([6u8; 32]);
 62    let mut source_xa: Source = Source::new([2u8; 32]);
 63    let mut source_xe: Source = Source::new([3u8; 32]);
 64
 65    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
 66
 67    let mut res: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_infos);
 68
 69    let k: u32 = source.next_u32();
 70
 71    let mut k_enc_prep: FheUintPrepared<Vec<u8>, u32, BE> =
 72        FheUintPrepared::<Vec<u8>, u32, BE>::alloc_from_infos(module, &ggsw_infos);
 73    k_enc_prep.encrypt_sk(
 74        module,
 75        k,
 76        sk_glwe_prep,
 77        &mut source_xa,
 78        &mut source_xe,
 79        scratch.borrow(),
 80    );
 81
 82    let digit = 5;
 83    let mask: u32 = (1 << digit) - 1;
 84
 85    // Starting bit
 86    let mut bit_start: usize = 0;
 87
 88    let mut data = vec![0i64; 1 << digit];
 89    data.iter_mut().enumerate().for_each(|(i, x)| *x = i as i64);
 90
 91    for _ in 0..32_usize.div_ceil(digit) {
 92        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_infos);
 93
 94        let mut cts_map: HashMap<usize, &mut GLWE<Vec<u8>>> = HashMap::new();
 95        let mut cts: Vec<GLWE<Vec<u8>>> = Vec::new();
 96
 97        for value in data.iter().take(1 << digit) {
 98            pt.encode_coeff_i64(*value, TorusPrecision(base2k.as_u32()), 0);
 99            let mut ct = GLWE::alloc_from_infos(&glwe_infos);
100            ct.encrypt_sk(
101                module,
102                &pt,
103                sk_glwe_prep,
104                &mut source_xa,
105                &mut source_xe,
106                scratch.borrow(),
107            );
108            cts.push(ct);
109        }
110
111        for (i, ct) in cts.iter_mut().enumerate() {
112            if i.is_multiple_of(3) {
113                cts_map.insert(i, ct);
114            }
115        }
116
117        // How many bits to take
118        let bit_size: usize = (32 - bit_start).min(digit);
119
120        module.glwe_blind_selection(
121            &mut res,
122            cts_map,
123            &k_enc_prep,
124            bit_start,
125            bit_size,
126            scratch.borrow(),
127        );
128
129        res.decrypt(module, &mut pt, sk_glwe_prep, scratch.borrow());
130
131        let idx = ((k >> bit_start) & mask) as usize;
132        if !idx.is_multiple_of(3) {
133            assert_eq!(0, pt.decode_coeff_i64(TorusPrecision(base2k.as_u32()), 0));
134        } else {
135            assert_eq!(
136                data[idx],
137                pt.decode_coeff_i64(TorusPrecision(base2k.as_u32()), 0)
138            );
139        }
140
141        bit_start += digit;
142
143        if bit_start >= 32 {
144            break;
145        }
146    }
147}