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