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