1use backend::hal::{
  2    api::{ScratchAvailable, TakeVecZnx, TakeVecZnxDft, ZnxView, ZnxViewMut, ZnxZero},
  3    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  4};
  5
  6use crate::{
  7    GLWECiphertext, GLWEKeyswitchFamily, GLWEToLWESwitchingKeyExec, Infos, LWECiphertext, LWESwitchingKeyExec,
  8    LWEToGLWESwitchingKeyExec, TakeGLWECt,
  9};
 10
 11impl LWECiphertext<Vec<u8>> {
 12    pub fn from_glwe_scratch_space<B: Backend>(
 13        module: &Module<B>,
 14        n: usize,
 15        basek: usize,
 16        k_lwe: usize,
 17        k_glwe: usize,
 18        k_ksk: usize,
 19        rank: usize,
 20    ) -> usize
 21    where
 22        Module<B>: GLWEKeyswitchFamily<B>,
 23    {
 24        GLWECiphertext::bytes_of(n, basek, k_lwe, 1)
 25            + GLWECiphertext::keyswitch_scratch_space(module, n, basek, k_lwe, k_glwe, k_ksk, 1, rank, 1)
 26    }
 27
 28    pub fn keyswitch_scratch_space<B: Backend>(
 29        module: &Module<B>,
 30        n: usize,
 31        basek: usize,
 32        k_lwe_out: usize,
 33        k_lwe_in: usize,
 34        k_ksk: usize,
 35    ) -> usize
 36    where
 37        Module<B>: GLWEKeyswitchFamily<B>,
 38    {
 39        GLWECiphertext::bytes_of(n, basek, k_lwe_out.max(k_lwe_in), 1)
 40            + GLWECiphertext::keyswitch_inplace_scratch_space(module, n, basek, k_lwe_out, k_ksk, 1, 1)
 41    }
 42}
 43
 44impl<DLwe: DataMut> LWECiphertext<DLwe> {
 45    pub fn sample_extract<DGlwe: DataRef>(&mut self, a: &GLWECiphertext<DGlwe>) {
 46        #[cfg(debug_assertions)]
 47        {
 48            assert!(self.n() <= a.n());
 49        }
 50
 51        let min_size: usize = self.size().min(a.size());
 52        let n: usize = self.n();
 53
 54        self.data.zero();
 55        (0..min_size).for_each(|i| {
 56            let data_lwe: &mut [i64] = self.data.at_mut(0, i);
 57            data_lwe[0] = a.data.at(0, i)[0];
 58            data_lwe[1..].copy_from_slice(&a.data.at(1, i)[..n]);
 59        });
 60    }
 61
 62    pub fn from_glwe<DGlwe, DKs, B: Backend>(
 63        &mut self,
 64        module: &Module<B>,
 65        a: &GLWECiphertext<DGlwe>,
 66        ks: &GLWEToLWESwitchingKeyExec<DKs, B>,
 67        scratch: &mut Scratch<B>,
 68    ) where
 69        DGlwe: DataRef,
 70        DKs: DataRef,
 71        Module<B>: GLWEKeyswitchFamily<B>,
 72        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
 73    {
 74        #[cfg(debug_assertions)]
 75        {
 76            assert_eq!(self.basek(), a.basek());
 77            assert_eq!(a.n(), ks.n());
 78        }
 79        let (mut tmp_glwe, scratch1) = scratch.take_glwe_ct(a.n(), a.basek(), self.k(), 1);
 80        tmp_glwe.keyswitch(module, a, &ks.0, scratch1);
 81        self.sample_extract(&tmp_glwe);
 82    }
 83
 84    pub fn keyswitch<A, DKs, B: Backend>(
 85        &mut self,
 86        module: &Module<B>,
 87        a: &LWECiphertext<A>,
 88        ksk: &LWESwitchingKeyExec<DKs, B>,
 89        scratch: &mut Scratch<B>,
 90    ) where
 91        A: DataRef,
 92        DKs: DataRef,
 93        Module<B>: GLWEKeyswitchFamily<B>,
 94        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
 95    {
 96        #[cfg(debug_assertions)]
 97        {
 98            assert!(self.n() <= module.n());
 99            assert!(a.n() <= module.n());
100            assert_eq!(self.basek(), a.basek());
101        }
102
103        let max_k: usize = self.k().max(a.k());
104        let basek: usize = self.basek();
105
106        let (mut glwe, scratch1) = scratch.take_glwe_ct(ksk.n(), basek, max_k, 1);
107        glwe.data.zero();
108
109        let n_lwe: usize = a.n();
110
111        (0..a.size()).for_each(|i| {
112            let data_lwe: &[i64] = a.data.at(0, i);
113            glwe.data.at_mut(0, i)[0] = data_lwe[0];
114            glwe.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]);
115        });
116
117        glwe.keyswitch_inplace(module, &ksk.0, scratch1);
118
119        self.sample_extract(&glwe);
120    }
121}
122
123impl GLWECiphertext<Vec<u8>> {
124    pub fn from_lwe_scratch_space<B: Backend>(
125        module: &Module<B>,
126        n: usize,
127        basek: usize,
128        k_lwe: usize,
129        k_glwe: usize,
130        k_ksk: usize,
131        rank: usize,
132    ) -> usize
133    where
134        Module<B>: GLWEKeyswitchFamily<B>,
135    {
136        GLWECiphertext::keyswitch_scratch_space(module, n, basek, k_glwe, k_lwe, k_ksk, 1, 1, rank)
137            + GLWECiphertext::bytes_of(n, basek, k_lwe, 1)
138    }
139}
140
141impl<D: DataMut> GLWECiphertext<D> {
142    pub fn from_lwe<DLwe, DKsk, B: Backend>(
143        &mut self,
144        module: &Module<B>,
145        lwe: &LWECiphertext<DLwe>,
146        ksk: &LWEToGLWESwitchingKeyExec<DKsk, B>,
147        scratch: &mut Scratch<B>,
148    ) where
149        DLwe: DataRef,
150        DKsk: DataRef,
151        Module<B>: GLWEKeyswitchFamily<B>,
152        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
153    {
154        #[cfg(debug_assertions)]
155        {
156            assert!(lwe.n() <= self.n());
157            assert_eq!(self.basek(), self.basek());
158        }
159
160        let (mut glwe, scratch1) = scratch.take_glwe_ct(ksk.n(), lwe.basek(), lwe.k(), 1);
161        glwe.data.zero();
162
163        let n_lwe: usize = lwe.n();
164
165        (0..lwe.size()).for_each(|i| {
166            let data_lwe: &[i64] = lwe.data.at(0, i);
167            glwe.data.at_mut(0, i)[0] = data_lwe[0];
168            glwe.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]);
169        });
170
171        self.keyswitch(module, &glwe, &ksk.0, scratch1);
172    }
173}