1use poulpy_hal::{
  2    api::{
  3        ScratchAvailable, SvpApplyDftToDftInplace, TakeVecZnx, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal,
  4        VecZnxAddScalarInplace, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf, VecZnxFillUniform, VecZnxIdftApplyConsume,
  5        VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace, VmpPMatAlloc, VmpPrepare,
  6    },
  7    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, ScalarZnxToRef, Scratch, ZnxView, ZnxViewMut},
  8    source::Source,
  9};
 10
 11use std::marker::PhantomData;
 12
 13use poulpy_core::{
 14    Distribution,
 15    layouts::{
 16        GGSW, GGSWInfos, LWESecret,
 17        compressed::GGSWCompressed,
 18        prepared::{GGSWPrepared, GLWESecretPrepared},
 19    },
 20};
 21
 22use crate::tfhe::blind_rotation::{
 23    BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyCompressed, BlindRotationKeyEncryptSk, BlindRotationKeyInfos,
 24    BlindRotationKeyPrepared, BlindRotationKeyPreparedAlloc, CGGI,
 25};
 26
 27impl BlindRotationKeyAlloc for BlindRotationKey<Vec<u8>, CGGI> {
 28    fn alloc<A>(infos: &A) -> Self
 29    where
 30        A: BlindRotationKeyInfos,
 31    {
 32        let mut data: Vec<GGSW<Vec<u8>>> = Vec::with_capacity(infos.n_lwe().into());
 33        for _ in 0..infos.n_lwe().as_usize() {
 34            data.push(GGSW::alloc_from_infos(infos));
 35        }
 36
 37        Self {
 38            keys: data,
 39            dist: Distribution::NONE,
 40            _phantom: PhantomData,
 41        }
 42    }
 43}
 44
 45impl BlindRotationKey<Vec<u8>, CGGI> {
 46    pub fn generate_from_sk_scratch_space<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
 47    where
 48        A: GGSWInfos,
 49        Module<B>: VecZnxNormalizeTmpBytes + VecZnxDftBytesOf,
 50    {
 51        GGSW::encrypt_sk_scratch_space(module, infos)
 52    }
 53}
 54
 55impl<D: DataMut, B: Backend> BlindRotationKeyEncryptSk<B> for BlindRotationKey<D, CGGI>
 56where
 57    Module<B>: VecZnxAddScalarInplace
 58        + VecZnxDftBytesOf
 59        + VecZnxBigNormalize<B>
 60        + VecZnxDftApply<B>
 61        + SvpApplyDftToDftInplace<B>
 62        + VecZnxIdftApplyConsume<B>
 63        + VecZnxNormalizeTmpBytes
 64        + VecZnxFillUniform
 65        + VecZnxSubInplace
 66        + VecZnxAddInplace
 67        + VecZnxNormalizeInplace<B>
 68        + VecZnxAddNormal
 69        + VecZnxNormalize<B>
 70        + VecZnxSub,
 71    Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
 72{
 73    fn encrypt_sk<DataSkGLWE, DataSkLWE>(
 74        &mut self,
 75        module: &Module<B>,
 76        sk_glwe: &GLWESecretPrepared<DataSkGLWE, B>,
 77        sk_lwe: &LWESecret<DataSkLWE>,
 78        source_xa: &mut Source,
 79        source_xe: &mut Source,
 80        scratch: &mut Scratch<B>,
 81    ) where
 82        DataSkGLWE: DataRef,
 83        DataSkLWE: DataRef,
 84    {
 85        #[cfg(debug_assertions)]
 86        {
 87            use poulpy_core::layouts::{GLWEInfos, LWEInfos};
 88
 89            assert_eq!(self.keys.len() as u32, sk_lwe.n());
 90            assert!(sk_glwe.n() <= module.n() as u32);
 91            assert_eq!(sk_glwe.rank(), self.rank());
 92            match sk_lwe.dist() {
 93                Distribution::BinaryBlock(_)
 94                | Distribution::BinaryFixed(_)
 95                | Distribution::BinaryProb(_)
 96                | Distribution::ZERO => {}
 97                _ => panic!(
 98                    "invalid GLWESecret distribution: must be BinaryBlock, BinaryFixed or BinaryProb (or ZERO for debugging)"
 99                ),
100            }
101        }
102
103        self.dist = sk_lwe.dist();
104
105        let mut pt: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(sk_glwe.n().into(), 1);
106        let sk_ref: ScalarZnx<&[u8]> = sk_lwe.data().to_ref();
107
108        self.keys.iter_mut().enumerate().for_each(|(i, ggsw)| {
109            pt.at_mut(0, 0)[0] = sk_ref.at(0, 0)[i];
110            ggsw.encrypt_sk(module, &pt, sk_glwe, source_xa, source_xe, scratch);
111        });
112    }
113}
114
115impl<B: Backend> BlindRotationKeyPreparedAlloc<B> for BlindRotationKeyPrepared<Vec<u8>, CGGI, B>
116where
117    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
118{
119    fn alloc<A>(module: &Module<B>, infos: &A) -> Self
120    where
121        A: BlindRotationKeyInfos,
122    {
123        let mut data: Vec<GGSWPrepared<Vec<u8>, B>> = Vec::with_capacity(infos.n_lwe().into());
124        (0..infos.n_lwe().as_usize()).for_each(|_| data.push(GGSWPrepared::alloc_from_infos(module, infos)));
125        Self {
126            data,
127            dist: Distribution::NONE,
128            x_pow_a: None,
129            _phantom: PhantomData,
130        }
131    }
132}
133
134impl BlindRotationKeyCompressed<Vec<u8>, CGGI> {
135    pub fn alloc<A>(infos: &A) -> Self
136    where
137        A: BlindRotationKeyInfos,
138    {
139        let mut data: Vec<GGSWCompressed<Vec<u8>>> = Vec::with_capacity(infos.n_lwe().into());
140        (0..infos.n_lwe().as_usize()).for_each(|_| data.push(GGSWCompressed::alloc_from_infos(infos)));
141        Self {
142            keys: data,
143            dist: Distribution::NONE,
144            _phantom: PhantomData,
145        }
146    }
147
148    pub fn generate_from_sk_scratch_space<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
149    where
150        A: GGSWInfos,
151        Module<B>: VecZnxNormalizeTmpBytes + VecZnxDftBytesOf,
152    {
153        GGSWCompressed::encrypt_sk_scratch_space(module, infos)
154    }
155}
156
157impl<D: DataMut> BlindRotationKeyCompressed<D, CGGI> {
158    #[allow(clippy::too_many_arguments)]
159    pub fn encrypt_sk<DataSkGLWE, DataSkLWE, B: Backend>(
160        &mut self,
161        module: &Module<B>,
162        sk_glwe: &GLWESecretPrepared<DataSkGLWE, B>,
163        sk_lwe: &LWESecret<DataSkLWE>,
164        seed_xa: [u8; 32],
165        source_xe: &mut Source,
166        scratch: &mut Scratch<B>,
167    ) where
168        DataSkGLWE: DataRef,
169        DataSkLWE: DataRef,
170        Module<B>: VecZnxAddScalarInplace
171            + VecZnxDftBytesOf
172            + VecZnxBigNormalize<B>
173            + VecZnxDftApply<B>
174            + SvpApplyDftToDftInplace<B>
175            + VecZnxIdftApplyConsume<B>
176            + VecZnxNormalizeTmpBytes
177            + VecZnxFillUniform
178            + VecZnxSubInplace
179            + VecZnxAddInplace
180            + VecZnxNormalizeInplace<B>
181            + VecZnxAddNormal
182            + VecZnxNormalize<B>
183            + VecZnxSub,
184        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
185    {
186        #[cfg(debug_assertions)]
187        {
188            use poulpy_core::layouts::{GLWEInfos, LWEInfos};
189
190            assert_eq!(self.n_lwe(), sk_lwe.n());
191            assert!(sk_glwe.n() <= module.n() as u32);
192            assert_eq!(sk_glwe.rank(), self.rank());
193            match sk_lwe.dist() {
194                Distribution::BinaryBlock(_)
195                | Distribution::BinaryFixed(_)
196                | Distribution::BinaryProb(_)
197                | Distribution::ZERO => {}
198                _ => panic!(
199                    "invalid GLWESecret distribution: must be BinaryBlock, BinaryFixed or BinaryProb (or ZERO for debugging)"
200                ),
201            }
202        }
203
204        self.dist = sk_lwe.dist();
205
206        let mut pt: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(sk_glwe.n().into(), 1);
207        let sk_ref: ScalarZnx<&[u8]> = sk_lwe.data().to_ref();
208
209        let mut source_xa: Source = Source::new(seed_xa);
210
211        self.keys.iter_mut().enumerate().for_each(|(i, ggsw)| {
212            pt.at_mut(0, 0)[0] = sk_ref.at(0, 0)[i];
213            ggsw.encrypt_sk(
214                module,
215                &pt,
216                sk_glwe,
217                source_xa.new_seed(),
218                source_xe,
219                scratch,
220            );
221        });
222    }
223}