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