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