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