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