1use poulpy_core::{
 2    GLWECopy, GLWERotate, ScratchTakeCore,
 3    layouts::{GGSW, GGSWInfos, GGSWToMut, GLWE, GLWEInfos, GLWEToMut},
 4};
 5use poulpy_hal::layouts::{Backend, Scratch};
 6
 7use crate::tfhe::bdd_arithmetic::{Cmux, GetGGSWBit, UnsignedInteger};
 8
 9pub trait GGSWBlindRotation<T: UnsignedInteger, BE: Backend>
10where
11    Self: GLWEBlindRotation<T, BE>,
12    Scratch<BE>: ScratchTakeCore<BE>,
13{
14    fn ggsw_blind_rotation_tmp_bytes<R, K>(&self, res_infos: &R, k_infos: &K) -> usize
15    where
16        R: GLWEInfos,
17        K: GGSWInfos,
18    {
19        self.glwe_blind_rotation_tmp_bytes(res_infos, k_infos)
20    }
21
22    fn ggsw_blind_rotation<R, K>(
23        &self,
24        res: &mut R,
25        k: &K,
26        bit_start: usize,
27        bit_size: usize,
28        bit_step: usize,
29        scratch: &mut Scratch<BE>,
30    ) where
31        R: GGSWToMut,
32        K: GetGGSWBit<T, BE>,
33        Scratch<BE>: ScratchTakeCore<BE>,
34    {
35        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
36
37        for row in 0..res.dnum().into() {
38            for col in 0..(res.rank() + 1).into() {
39                self.glwe_blind_rotation(
40                    &mut res.at_mut(row, col),
41                    k,
42                    bit_start,
43                    bit_size,
44                    bit_step,
45                    scratch,
46                );
47            }
48        }
49    }
50}
51
52pub trait GLWEBlindRotation<T: UnsignedInteger, BE: Backend>
53where
54    Self: GLWECopy + GLWERotate<BE> + Cmux<BE>,
55    Scratch<BE>: ScratchTakeCore<BE>,
56{
57    fn glwe_blind_rotation_tmp_bytes<R, K>(&self, res_infos: &R, k_infos: &K) -> usize
58    where
59        R: GLWEInfos,
60        K: GGSWInfos,
61    {
62        self.cmux_tmp_bytes(res_infos, res_infos, k_infos) + GLWE::bytes_of_from_infos(res_infos)
63    }
64
65    /// Homomorphic multiplication of res by X^{k[bit_start..bit_start + bit_size] * bit_step}.
66    fn glwe_blind_rotation<R, K>(
67        &self,
68        res: &mut R,
69        k: &K,
70        bit_start: usize,
71        bit_size: usize,
72        bit_step: usize,
73        scratch: &mut Scratch<BE>,
74    ) where
75        R: GLWEToMut,
76        K: GetGGSWBit<T, BE>,
77        Scratch<BE>: ScratchTakeCore<BE>,
78    {
79        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
80
81        let (mut tmp_res, scratch_1) = scratch.take_glwe(res);
82
83        self.glwe_copy(&mut tmp_res, res);
84
85        for i in 1..bit_size {
86            // res' = res * X^2^(i * bit_step)
87            self.glwe_rotate(1 << (i + bit_step), &mut tmp_res, res);
88
89            // res = (res - res') * GGSW(b[i]) + res'
90            self.cmux_inplace(res, &tmp_res, &k.get_bit(i + bit_start), scratch_1);
91        }
92    }
93}