1use poulpy_core::{
  2    GLWECopy, GLWERotate, ScratchTakeCore,
  3    layouts::{GGSW, GGSWInfos, GGSWToMut, GGSWToRef, GLWE, GLWEInfos, GLWEToMut, GLWEToRef},
  4};
  5use poulpy_hal::layouts::{Backend, Module, Scratch};
  6
  7use crate::tfhe::bdd_arithmetic::{Cmux, GetGGSWBit, UnsignedInteger};
  8
  9impl<T: UnsignedInteger, BE: Backend> GGSWBlindRotation<T, BE> for Module<BE>
 10where
 11    Self: GLWEBlindRotation<T, BE>,
 12    Scratch<BE>: ScratchTakeCore<BE>,
 13{
 14}
 15
 16pub trait GGSWBlindRotation<T: UnsignedInteger, BE: Backend>
 17where
 18    Self: GLWEBlindRotation<T, BE>,
 19    Scratch<BE>: ScratchTakeCore<BE>,
 20{
 21    fn ggsw_blind_rotation_tmp_bytes<R, K>(&self, res_infos: &R, k_infos: &K) -> usize
 22    where
 23        R: GLWEInfos,
 24        K: GGSWInfos,
 25    {
 26        self.glwe_blind_rotation_tmp_bytes(res_infos, k_infos)
 27    }
 28
 29    fn ggsw_blind_rotation<R, G, K>(
 30        &self,
 31        res: &mut R,
 32        test_ggsw: &G,
 33        k: &K,
 34        bit_start: usize,
 35        bit_size: usize,
 36        bit_step: usize,
 37        scratch: &mut Scratch<BE>,
 38    ) where
 39        R: GGSWToMut,
 40        G: GGSWToRef,
 41        K: GetGGSWBit<T, BE>,
 42        Scratch<BE>: ScratchTakeCore<BE>,
 43    {
 44        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
 45        let test_ggsw: &GGSW<&[u8]> = &test_ggsw.to_ref();
 46
 47        for row in 0..res.dnum().into() {
 48            for col in 0..(res.rank() + 1).into() {
 49                self.glwe_blind_rotation(
 50                    &mut res.at_mut(row, col),
 51                    &test_ggsw.at(row, col),
 52                    k,
 53                    bit_start,
 54                    bit_size,
 55                    bit_step,
 56                    scratch,
 57                );
 58            }
 59        }
 60    }
 61}
 62
 63impl<T: UnsignedInteger, BE: Backend> GLWEBlindRotation<T, BE> for Module<BE>
 64where
 65    Self: GLWECopy + GLWERotate<BE> + Cmux<BE>,
 66    Scratch<BE>: ScratchTakeCore<BE>,
 67{
 68}
 69
 70pub trait GLWEBlindRotation<T: UnsignedInteger, BE: Backend>
 71where
 72    Self: GLWECopy + GLWERotate<BE> + Cmux<BE>,
 73    Scratch<BE>: ScratchTakeCore<BE>,
 74{
 75    fn glwe_blind_rotation_tmp_bytes<R, K>(&self, res_infos: &R, k_infos: &K) -> usize
 76    where
 77        R: GLWEInfos,
 78        K: GGSWInfos,
 79    {
 80        self.cmux_tmp_bytes(res_infos, res_infos, k_infos) + GLWE::bytes_of_from_infos(res_infos)
 81    }
 82
 83    /// Homomorphic multiplication of res by X^{k[bit_start..bit_start + bit_size] * bit_step}.
 84    fn glwe_blind_rotation<R, G, K>(
 85        &self,
 86        res: &mut R,
 87        test_glwe: &G,
 88        k: &K,
 89        bit_start: usize,
 90        bit_size: usize,
 91        bit_step: usize,
 92        scratch: &mut Scratch<BE>,
 93    ) where
 94        R: GLWEToMut,
 95        G: GLWEToRef,
 96        K: GetGGSWBit<T, BE>,
 97        Scratch<BE>: ScratchTakeCore<BE>,
 98    {
 99        assert!(bit_start + bit_size <= T::WORD_SIZE);
100
101        let mut res: GLWE<&mut [u8]> = res.to_mut();
102
103        let (mut tmp_res, scratch_1) = scratch.take_glwe(&res);
104
105        // res <- test_glwe
106        self.glwe_copy(&mut res, test_glwe);
107
108        // a_is_res = true  => (a, b) = (&mut res, &mut tmp_res)
109        // a_is_res = false => (a, b) = (&mut tmp_res, &mut res)
110        let mut a_is_res: bool = true;
111
112        for i in 0..bit_size {
113            let (a, b) = if a_is_res {
114                (&mut res, &mut tmp_res)
115            } else {
116                (&mut tmp_res, &mut res)
117            };
118
119            // a <- a ; b <- a * X^{-2^{i + bit_step}}
120            self.glwe_rotate(-1 << (i + bit_step), b, a);
121
122            // b <- (b - a) * GGSW(b[i]) + a
123            self.cmux_inplace(b, a, &k.get_bit(i + bit_start), scratch_1);
124
125            // ping-pong roles for next iter
126            a_is_res = !a_is_res;
127        }
128
129        // Ensure the final value ends up in `res`
130        if !a_is_res {
131            self.glwe_copy(&mut res, &tmp_res);
132        }
133    }
134}