1use poulpy_core::{
  2    GLWECopy, GLWERotate, ScratchTakeCore,
  3    layouts::{GGSW, GGSWInfos, GGSWToMut, GGSWToRef, GLWE, GLWEInfos, GLWEToMut, GLWEToRef, LWEInfos},
  4};
  5use poulpy_hal::{
  6    api::{VecZnxAddScalarInplace, VecZnxNormalizeInplace},
  7    layouts::{Backend, Module, ScalarZnx, ScalarZnxToRef, Scratch, ZnxZero},
  8};
  9
 10use crate::tfhe::bdd_arithmetic::{Cmux, GetGGSWBit, UnsignedInteger};
 11
 12impl<T: UnsignedInteger, BE: Backend> GGSWBlindRotation<T, BE> for Module<BE>
 13where
 14    Self: GLWEBlindRotation<T, BE> + VecZnxAddScalarInplace + VecZnxNormalizeInplace<BE>,
 15    Scratch<BE>: ScratchTakeCore<BE>,
 16{
 17}
 18
 19pub trait GGSWBlindRotation<T: UnsignedInteger, BE: Backend>
 20where
 21    Self: GLWEBlindRotation<T, BE> + VecZnxAddScalarInplace + VecZnxNormalizeInplace<BE>,
 22{
 23    fn ggsw_to_ggsw_blind_rotation_tmp_bytes<R, K>(&self, res_infos: &R, k_infos: &K) -> usize
 24    where
 25        R: GLWEInfos,
 26        K: GGSWInfos,
 27    {
 28        self.glwe_blind_rotation_tmp_bytes(res_infos, k_infos)
 29    }
 30
 31    #[allow(clippy::too_many_arguments)]
 32    /// res <- res * X^{((k>>bit_rsh) % 2^bit_mask) << bit_lsh}.
 33    fn ggsw_blind_rotation_inplace<R, K>(
 34        &self,
 35        res: &mut R,
 36        fhe_uint: &K,
 37        sign: bool,
 38        bit_rsh: usize,
 39        bit_mask: usize,
 40        bit_lsh: usize,
 41        scratch: &mut Scratch<BE>,
 42    ) where
 43        R: GGSWToMut,
 44        K: GetGGSWBit<BE>,
 45        Scratch<BE>: ScratchTakeCore<BE>,
 46    {
 47        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
 48
 49        for col in 0..(res.rank() + 1).into() {
 50            for row in 0..res.dnum().into() {
 51                self.glwe_blind_rotation_inplace(
 52                    &mut res.at_mut(row, col),
 53                    fhe_uint,
 54                    sign,
 55                    bit_rsh,
 56                    bit_mask,
 57                    bit_lsh,
 58                    scratch,
 59                );
 60            }
 61        }
 62    }
 63
 64    #[allow(clippy::too_many_arguments)]
 65    /// res <- a * X^{((k>>bit_rsh) % 2^bit_mask) << bit_lsh}.
 66    fn ggsw_blind_rotation<R, A, K>(
 67        &self,
 68        res: &mut R,
 69        a: &A,
 70        fhe_uint: &K,
 71        sign: bool,
 72        bit_rsh: usize,
 73        bit_mask: usize,
 74        bit_lsh: usize,
 75        scratch: &mut Scratch<BE>,
 76    ) where
 77        R: GGSWToMut,
 78        A: GGSWToRef,
 79        K: GetGGSWBit<BE>,
 80        Scratch<BE>: ScratchTakeCore<BE>,
 81    {
 82        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
 83        let a: &GGSW<&[u8]> = &a.to_ref();
 84
 85        assert!(res.dnum() <= a.dnum());
 86        assert_eq!(res.dsize(), a.dsize());
 87
 88        for col in 0..(res.rank() + 1).into() {
 89            for row in 0..res.dnum().into() {
 90                self.glwe_blind_rotation(
 91                    &mut res.at_mut(row, col),
 92                    &a.at(row, col),
 93                    fhe_uint,
 94                    sign,
 95                    bit_rsh,
 96                    bit_mask,
 97                    bit_lsh,
 98                    scratch,
 99                );
100            }
101        }
102    }
103
104    fn scalar_to_ggsw_blind_rotation_tmp_bytes<R, K>(&self, res_infos: &R, k_infos: &K) -> usize
105    where
106        R: GLWEInfos,
107        K: GGSWInfos,
108    {
109        self.glwe_blind_rotation_tmp_bytes(res_infos, k_infos) + GLWE::bytes_of_from_infos(res_infos)
110    }
111
112    #[allow(clippy::too_many_arguments)]
113    fn scalar_to_ggsw_blind_rotation<R, S, K>(
114        &self,
115        res: &mut R,
116        test_vector: &S,
117        fhe_uint: &K,
118        sign: bool,
119        bit_rsh: usize,
120        bit_mask: usize,
121        bit_lsh: usize,
122        scratch: &mut Scratch<BE>,
123    ) where
124        R: GGSWToMut,
125        S: ScalarZnxToRef,
126        K: GetGGSWBit<BE>,
127        Scratch<BE>: ScratchTakeCore<BE>,
128    {
129        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
130        let test_vector: &ScalarZnx<&[u8]> = &test_vector.to_ref();
131
132        let base2k: usize = res.base2k().into();
133        let dsize: usize = res.dsize().into();
134
135        let (mut tmp_glwe, scratch_1) = scratch.take_glwe(res);
136
137        for col in 0..(res.rank() + 1).into() {
138            for row in 0..res.dnum().into() {
139                tmp_glwe.data_mut().zero();
140                self.vec_znx_add_scalar_inplace(
141                    tmp_glwe.data_mut(),
142                    col,
143                    (dsize - 1) + row * dsize,
144                    test_vector,
145                    0,
146                );
147                self.vec_znx_normalize_inplace(base2k, tmp_glwe.data_mut(), col, scratch_1);
148
149                self.glwe_blind_rotation(
150                    &mut res.at_mut(row, col),
151                    &tmp_glwe,
152                    fhe_uint,
153                    sign,
154                    bit_rsh,
155                    bit_mask,
156                    bit_lsh,
157                    scratch_1,
158                );
159            }
160        }
161    }
162}
163
164impl<T: UnsignedInteger, BE: Backend> GLWEBlindRotation<T, BE> for Module<BE>
165where
166    Self: GLWECopy + GLWERotate<BE> + Cmux<BE>,
167    Scratch<BE>: ScratchTakeCore<BE>,
168{
169}
170
171pub trait GLWEBlindRotation<T: UnsignedInteger, BE: Backend>
172where
173    Self: GLWECopy + GLWERotate<BE> + Cmux<BE>,
174{
175    fn glwe_blind_rotation_tmp_bytes<R, K>(&self, res_infos: &R, k_infos: &K) -> usize
176    where
177        R: GLWEInfos,
178        K: GGSWInfos,
179    {
180        self.cmux_tmp_bytes(res_infos, res_infos, k_infos) + GLWE::bytes_of_from_infos(res_infos)
181    }
182
183    #[allow(clippy::too_many_arguments)]
184    fn glwe_blind_rotation_inplace<R, K>(
185        &self,
186        res: &mut R,
187        fhe_uint: &K,
188        sign: bool,
189        bit_rsh: usize,
190        bit_mask: usize,
191        bit_lsh: usize,
192        scratch: &mut Scratch<BE>,
193    ) where
194        R: GLWEToMut,
195        K: GetGGSWBit<BE>,
196        Scratch<BE>: ScratchTakeCore<BE>,
197    {
198        assert!(bit_rsh + bit_mask <= T::WORD_SIZE);
199
200        let mut res: GLWE<&mut [u8]> = res.to_mut();
201
202        let (mut tmp_res, scratch_1) = scratch.take_glwe(&res);
203
204        // a_is_res = true  => (a, b) = (&mut res, &mut tmp_res)
205        // a_is_res = false => (a, b) = (&mut tmp_res, &mut res)
206        let mut a_is_res: bool = true;
207
208        for i in 0..bit_mask {
209            let (a, b) = if a_is_res {
210                (&mut res, &mut tmp_res)
211            } else {
212                (&mut tmp_res, &mut res)
213            };
214
215            // a <- a ; b <- a * X^{-2^{i + bit_lsh}}
216            match sign {
217                true => self.glwe_rotate(1 << (i + bit_lsh), b, a),
218                false => self.glwe_rotate(-1 << (i + bit_lsh), b, a),
219            }
220
221            // b <- (b - a) * GGSW(b[i]) + a
222            self.cmux_inplace(b, a, &fhe_uint.get_bit(i + bit_rsh), scratch_1);
223
224            // ping-pong roles for next iter
225            a_is_res = !a_is_res;
226        }
227
228        // Ensure the final value ends up in `res`
229        if !a_is_res {
230            self.glwe_copy(&mut res, &tmp_res);
231        }
232    }
233
234    #[allow(clippy::too_many_arguments)]
235    /// res <- a * X^{sign * ((k>>bit_rsh) % 2^bit_mask) << bit_lsh}.
236    fn glwe_blind_rotation<R, A, K>(
237        &self,
238        res: &mut R,
239        a: &A,
240        fhe_uint: &K,
241        sign: bool,
242        bit_rsh: usize,
243        bit_mask: usize,
244        bit_lsh: usize,
245        scratch: &mut Scratch<BE>,
246    ) where
247        R: GLWEToMut,
248        A: GLWEToRef,
249        K: GetGGSWBit<BE>,
250        Scratch<BE>: ScratchTakeCore<BE>,
251    {
252        self.glwe_copy(res, a);
253        self.glwe_blind_rotation_inplace(res, fhe_uint, sign, bit_rsh, bit_mask, bit_lsh, scratch);
254    }
255}