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