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