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