1use poulpy_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,
  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        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>, basek: usize, k: usize, rank: usize) -> usize
 42    where
 43        Module<B>: VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes,
 44    {
 45        GGSWCiphertext::encrypt_sk_scratch_space(module, 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        scratch: &mut Scratch<B>,
 75    ) where
 76        DataSkGLWE: DataRef,
 77        DataSkLWE: DataRef,
 78    {
 79        #[cfg(debug_assertions)]
 80        {
 81            assert_eq!(self.keys.len(), sk_lwe.n());
 82            assert!(sk_glwe.n() <= module.n());
 83            assert_eq!(sk_glwe.rank(), self.keys[0].rank());
 84            match sk_lwe.dist() {
 85                Distribution::BinaryBlock(_)
 86                | Distribution::BinaryFixed(_)
 87                | Distribution::BinaryProb(_)
 88                | Distribution::ZERO => {}
 89                _ => panic!(
 90                    "invalid GLWESecret distribution: must be BinaryBlock, BinaryFixed or BinaryProb (or ZERO for debugging)"
 91                ),
 92            }
 93        }
 94
 95        self.dist = sk_lwe.dist();
 96
 97        let mut pt: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(sk_glwe.n(), 1);
 98        let sk_ref: ScalarZnx<&[u8]> = sk_lwe.data().to_ref();
 99
100        self.keys.iter_mut().enumerate().for_each(|(i, ggsw)| {
101            pt.at_mut(0, 0)[0] = sk_ref.at(0, 0)[i];
102            ggsw.encrypt_sk(module, &pt, sk_glwe, source_xa, source_xe, scratch);
103        });
104    }
105}
106
107impl<B: Backend> BlindRotationKeyPreparedAlloc<B> for BlindRotationKeyPrepared<Vec<u8>, CGGI, B>
108where
109    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
110{
111    fn alloc(module: &Module<B>, n_lwe: usize, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
112        let mut data: Vec<GGSWCiphertextPrepared<Vec<u8>, B>> = Vec::with_capacity(n_lwe);
113        (0..n_lwe).for_each(|_| {
114            data.push(GGSWCiphertextPrepared::alloc(
115                module, basek, k, rows, 1, rank,
116            ))
117        });
118        Self {
119            data,
120            dist: Distribution::NONE,
121            x_pow_a: None,
122            _phantom: PhantomData,
123        }
124    }
125}
126
127impl BlindRotationKeyCompressed<Vec<u8>, CGGI> {
128    pub fn alloc(n_gglwe: usize, n_lwe: usize, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
129        let mut data: Vec<GGSWCiphertextCompressed<Vec<u8>>> = Vec::with_capacity(n_lwe);
130        (0..n_lwe).for_each(|_| {
131            data.push(GGSWCiphertextCompressed::alloc(
132                n_gglwe, basek, k, rows, 1, rank,
133            ))
134        });
135        Self {
136            keys: data,
137            dist: Distribution::NONE,
138            _phantom: PhantomData,
139        }
140    }
141
142    pub fn generate_from_sk_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize
143    where
144        Module<B>: VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes,
145    {
146        GGSWCiphertextCompressed::encrypt_sk_scratch_space(module, basek, k, rank)
147    }
148}
149
150impl<D: DataMut> BlindRotationKeyCompressed<D, CGGI> {
151    #[allow(clippy::too_many_arguments)]
152    pub fn encrypt_sk<DataSkGLWE, DataSkLWE, B: Backend>(
153        &mut self,
154        module: &Module<B>,
155        sk_glwe: &GLWESecretPrepared<DataSkGLWE, B>,
156        sk_lwe: &LWESecret<DataSkLWE>,
157        seed_xa: [u8; 32],
158        source_xe: &mut Source,
159        scratch: &mut Scratch<B>,
160    ) where
161        DataSkGLWE: DataRef,
162        DataSkLWE: DataRef,
163        Module<B>: VecZnxAddScalarInplace
164            + VecZnxDftAllocBytes
165            + VecZnxBigNormalize<B>
166            + VecZnxDftFromVecZnx<B>
167            + SvpApplyInplace<B>
168            + VecZnxDftToVecZnxBigConsume<B>
169            + VecZnxNormalizeTmpBytes
170            + VecZnxFillUniform
171            + VecZnxSubABInplace
172            + VecZnxAddInplace
173            + VecZnxNormalizeInplace<B>
174            + VecZnxAddNormal
175            + VecZnxNormalize<B>
176            + VecZnxSub,
177        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
178    {
179        #[cfg(debug_assertions)]
180        {
181            assert_eq!(self.keys.len(), sk_lwe.n());
182            assert!(sk_glwe.n() <= module.n());
183            assert_eq!(sk_glwe.rank(), self.keys[0].rank());
184            match sk_lwe.dist() {
185                Distribution::BinaryBlock(_)
186                | Distribution::BinaryFixed(_)
187                | Distribution::BinaryProb(_)
188                | Distribution::ZERO => {}
189                _ => panic!(
190                    "invalid GLWESecret distribution: must be BinaryBlock, BinaryFixed or BinaryProb (or ZERO for debugging)"
191                ),
192            }
193        }
194
195        self.dist = sk_lwe.dist();
196
197        let mut pt: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(sk_glwe.n(), 1);
198        let sk_ref: ScalarZnx<&[u8]> = sk_lwe.data().to_ref();
199
200        let mut source_xa: Source = Source::new(seed_xa);
201
202        self.keys.iter_mut().enumerate().for_each(|(i, ggsw)| {
203            pt.at_mut(0, 0)[0] = sk_ref.at(0, 0)[i];
204            ggsw.encrypt_sk(
205                module,
206                &pt,
207                sk_glwe,
208                source_xa.new_seed(),
209                source_xe,
210                scratch,
211            );
212        });
213    }
214}