1mod cggi;
  2
  3pub use cggi::*;
  4
  5use itertools::izip;
  6use poulpy_core::{
  7    ScratchTakeCore,
  8    layouts::{GGSWInfos, GLWE, GLWEInfos, LWE, LWEInfos},
  9};
 10use poulpy_hal::layouts::{Backend, DataMut, DataRef, Scratch, ZnxView};
 11
 12use crate::tfhe::blind_rotation::{BlindRotationKeyInfos, BlindRotationKeyPrepared, LookUpTableRotationDirection, LookupTable};
 13
 14pub trait BlindRotationAlgo: Sync {}
 15
 16pub trait BlindRotationExecute<BRA: BlindRotationAlgo, BE: Backend> {
 17    fn blind_rotation_execute_tmp_bytes<G, B>(
 18        &self,
 19        block_size: usize,
 20        extension_factor: usize,
 21        glwe_infos: &G,
 22        brk_infos: &B,
 23    ) -> usize
 24    where
 25        G: GLWEInfos,
 26        B: GGSWInfos;
 27
 28    fn blind_rotation_execute<DR, DL, DB>(
 29        &self,
 30        res: &mut GLWE<DR>,
 31        lwe: &LWE<DL>,
 32        lut: &LookupTable,
 33        brk: &BlindRotationKeyPrepared<DB, BRA, BE>,
 34        scratch: &mut Scratch<BE>,
 35    ) where
 36        DR: DataMut,
 37        DL: DataRef,
 38        DB: DataRef;
 39}
 40
 41impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> BlindRotationKeyPrepared<D, BRA, BE>
 42where
 43    Scratch<BE>: ScratchTakeCore<BE>,
 44{
 45    pub fn execute<DR: DataMut, DI: DataRef, M>(
 46        &self,
 47        module: &M,
 48        res: &mut GLWE<DR>,
 49        lwe: &LWE<DI>,
 50        lut: &LookupTable,
 51        scratch: &mut Scratch<BE>,
 52    ) where
 53        M: BlindRotationExecute<BRA, BE>,
 54    {
 55        module.blind_rotation_execute(res, lwe, lut, self, scratch);
 56    }
 57}
 58
 59impl<BE: Backend, BRA: BlindRotationAlgo> BlindRotationKeyPrepared<Vec<u8>, BRA, BE> {
 60    pub fn execute_tmp_bytes<A, B, M>(
 61        module: &M,
 62        block_size: usize,
 63        extension_factor: usize,
 64        glwe_infos: &A,
 65        brk_infos: &B,
 66    ) -> usize
 67    where
 68        A: GLWEInfos,
 69        B: BlindRotationKeyInfos,
 70        M: BlindRotationExecute<BRA, BE>,
 71    {
 72        module.blind_rotation_execute_tmp_bytes(block_size, extension_factor, glwe_infos, brk_infos)
 73    }
 74}
 75
 76pub fn mod_switch_2n(n: usize, res: &mut [i64], lwe: &LWE<&[u8]>, rot_dir: LookUpTableRotationDirection) {
 77    let base2k: usize = lwe.base2k().into();
 78
 79    let log2n: usize = usize::BITS as usize - (n - 1).leading_zeros() as usize + 1;
 80
 81    res.copy_from_slice(lwe.data().at(0, 0));
 82
 83    match rot_dir {
 84        LookUpTableRotationDirection::Left => {
 85            res.iter_mut().for_each(|x| *x = -*x);
 86        }
 87        LookUpTableRotationDirection::Right => {}
 88    }
 89
 90    if base2k > log2n {
 91        let diff: usize = base2k - (log2n - 1); // additional -1 because we map to [-N/2, N/2) instead of [0, N)
 92        res.iter_mut().for_each(|x| {
 93            *x = div_round_by_pow2(x, diff);
 94        })
 95    } else {
 96        let rem: usize = base2k - (log2n % base2k);
 97        let size: usize = log2n.div_ceil(base2k);
 98        (1..size).for_each(|i| {
 99            if i == size - 1 && rem != base2k {
100                let k_rem: usize = base2k - rem;
101                izip!(lwe.data().at(0, i).iter(), res.iter_mut()).for_each(|(x, y)| {
102                    *y = (*y << k_rem) + (x >> rem);
103                });
104            } else {
105                izip!(lwe.data().at(0, i).iter(), res.iter_mut()).for_each(|(x, y)| {
106                    *y = (*y << base2k) + x;
107                });
108            }
109        })
110    }
111}
112
113#[inline(always)]
114fn div_round_by_pow2(x: &i64, k: usize) -> i64 {
115    (x + (1 << (k - 1))) >> k
116}