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