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_to_glwe_blind_rotation_tmp_bytes(res_infos, k_infos)
 29    }
 30
 31    #[allow(clippy::too_many_arguments)]
 32    /// res <- a * X^{((k>>bit_rsh) % 2^bit_mask) << bit_lsh}.
 33    fn ggsw_to_ggsw_blind_rotation<R, A, K>(
 34        &self,
 35        res: &mut R,
 36        a: &A,
 37        k: &K,
 38        sign: bool,
 39        bit_rsh: usize,
 40        bit_mask: usize,
 41        bit_lsh: usize,
 42        scratch: &mut Scratch<BE>,
 43    ) where
 44        R: GGSWToMut,
 45        A: GGSWToRef,
 46        K: GetGGSWBit<T, BE>,
 47        Scratch<BE>: ScratchTakeCore<BE>,
 48    {
 49        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
 50        let a: &GGSW<&[u8]> = &a.to_ref();
 51
 52        assert!(res.dnum() <= a.dnum());
 53        assert_eq!(res.dsize(), a.dsize());
 54
 55        for col in 0..(res.rank() + 1).into() {
 56            for row in 0..res.dnum().into() {
 57                self.glwe_to_glwe_blind_rotation(
 58                    &mut res.at_mut(row, col),
 59                    &a.at(row, col),
 60                    k,
 61                    sign,
 62                    bit_rsh,
 63                    bit_mask,
 64                    bit_lsh,
 65                    scratch,
 66                );
 67            }
 68        }
 69    }
 70
 71    fn scalar_to_ggsw_blind_rotation_tmp_bytes<R, K>(&self, res_infos: &R, k_infos: &K) -> usize
 72    where
 73        R: GLWEInfos,
 74        K: GGSWInfos,
 75    {
 76        self.glwe_to_glwe_blind_rotation_tmp_bytes(res_infos, k_infos) + GLWE::bytes_of_from_infos(res_infos)
 77    }
 78
 79    #[allow(clippy::too_many_arguments)]
 80    fn scalar_to_ggsw_blind_rotation<R, S, K>(
 81        &self,
 82        res: &mut R,
 83        test_vector: &S,
 84        k: &K,
 85        sign: bool,
 86        bit_rsh: usize,
 87        bit_mask: usize,
 88        bit_lsh: usize,
 89        scratch: &mut Scratch<BE>,
 90    ) where
 91        R: GGSWToMut,
 92        S: ScalarZnxToRef,
 93        K: GetGGSWBit<T, BE>,
 94        Scratch<BE>: ScratchTakeCore<BE>,
 95    {
 96        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
 97        let test_vector: &ScalarZnx<&[u8]> = &test_vector.to_ref();
 98
 99        let base2k: usize = res.base2k().into();
100        let dsize: usize = res.dsize().into();
101
102        let (mut tmp_glwe, scratch_1) = scratch.take_glwe(res);
103
104        for col in 0..(res.rank() + 1).into() {
105            for row in 0..res.dnum().into() {
106                tmp_glwe.data_mut().zero();
107                self.vec_znx_add_scalar_inplace(
108                    tmp_glwe.data_mut(),
109                    col,
110                    (dsize - 1) + row * dsize,
111                    test_vector,
112                    0,
113                );
114                self.vec_znx_normalize_inplace(base2k, tmp_glwe.data_mut(), col, scratch_1);
115
116                self.glwe_to_glwe_blind_rotation(
117                    &mut res.at_mut(row, col),
118                    &tmp_glwe,
119                    k,
120                    sign,
121                    bit_rsh,
122                    bit_mask,
123                    bit_lsh,
124                    scratch_1,
125                );
126            }
127        }
128    }
129}
130
131impl<T: UnsignedInteger, BE: Backend> GLWEBlindRotation<T, BE> for Module<BE>
132where
133    Self: GLWECopy + GLWERotate<BE> + Cmux<BE>,
134    Scratch<BE>: ScratchTakeCore<BE>,
135{
136}
137
138pub trait GLWEBlindRotation<T: UnsignedInteger, BE: Backend>
139where
140    Self: GLWECopy + GLWERotate<BE> + Cmux<BE>,
141{
142    fn glwe_to_glwe_blind_rotation_tmp_bytes<R, K>(&self, res_infos: &R, k_infos: &K) -> usize
143    where
144        R: GLWEInfos,
145        K: GGSWInfos,
146    {
147        self.cmux_tmp_bytes(res_infos, res_infos, k_infos) + GLWE::bytes_of_from_infos(res_infos)
148    }
149
150    #[allow(clippy::too_many_arguments)]
151    /// res <- a * X^{sign * ((k>>bit_rsh) % 2^bit_mask) << bit_lsh}.
152    fn glwe_to_glwe_blind_rotation<R, A, K>(
153        &self,
154        res: &mut R,
155        a: &A,
156        k: &K,
157        sign: bool,
158        bit_rsh: usize,
159        bit_mask: usize,
160        bit_lsh: usize,
161        scratch: &mut Scratch<BE>,
162    ) where
163        R: GLWEToMut,
164        A: GLWEToRef,
165        K: GetGGSWBit<T, BE>,
166        Scratch<BE>: ScratchTakeCore<BE>,
167    {
168        assert!(bit_rsh + bit_mask <= T::WORD_SIZE);
169
170        let mut res: GLWE<&mut [u8]> = res.to_mut();
171        let a: &GLWE<&[u8]> = &a.to_ref();
172
173        let (mut tmp_res, scratch_1) = scratch.take_glwe(&res);
174
175        // a <- a ; b <- a * X^{-2^{i + bit_lsh}}
176        match sign {
177            true => self.glwe_rotate(1 << bit_lsh, &mut res, a),
178            false => self.glwe_rotate(-1 << bit_lsh, &mut res, a),
179        }
180
181        // b <- (b - a) * GGSW(b[i]) + a
182        self.cmux_inplace(&mut res, a, &k.get_bit(bit_rsh), scratch_1);
183
184        // a_is_res = true  => (a, b) = (&mut res, &mut tmp_res)
185        // a_is_res = false => (a, b) = (&mut tmp_res, &mut res)
186        let mut a_is_res: bool = true;
187
188        for i in 1..bit_mask {
189            let (a, b) = if a_is_res {
190                (&mut res, &mut tmp_res)
191            } else {
192                (&mut tmp_res, &mut res)
193            };
194
195            // a <- a ; b <- a * X^{-2^{i + bit_lsh}}
196            match sign {
197                true => self.glwe_rotate(1 << (i + bit_lsh), b, a),
198                false => self.glwe_rotate(-1 << (i + bit_lsh), b, a),
199            }
200
201            // b <- (b - a) * GGSW(b[i]) + a
202            self.cmux_inplace(b, a, &k.get_bit(i + bit_rsh), scratch_1);
203
204            // ping-pong roles for next iter
205            a_is_res = !a_is_res;
206        }
207
208        // Ensure the final value ends up in `res`
209        if !a_is_res {
210            self.glwe_copy(&mut res, &tmp_res);
211        }
212    }
213}