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<DLwe: DataMut> LWECiphertext<DLwe> {
 12    pub fn sample_extract<DGlwe: DataRef>(&mut self, a: &GLWECiphertext<DGlwe>) {
 13        #[cfg(debug_assertions)]
 14        {
 15            assert!(self.n() <= a.n());
 16        }
 17
 18        let min_size: usize = self.size().min(a.size());
 19        let n: usize = self.n();
 20
 21        self.data.zero();
 22        (0..min_size).for_each(|i| {
 23            let data_lwe: &mut [i64] = self.data.at_mut(0, i);
 24            data_lwe[0] = a.data.at(0, i)[0];
 25            data_lwe[1..].copy_from_slice(&a.data.at(1, i)[..n]);
 26        });
 27    }
 28
 29    pub fn from_glwe<DGlwe, DKs, B: Backend>(
 30        &mut self,
 31        module: &Module<B>,
 32        a: &GLWECiphertext<DGlwe>,
 33        ks: &GLWEToLWESwitchingKeyExec<DKs, B>,
 34        scratch: &mut Scratch<B>,
 35    ) where
 36        DGlwe: DataRef,
 37        DKs: DataRef,
 38        Module<B>: GLWEKeyswitchFamily<B>,
 39        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
 40    {
 41        #[cfg(debug_assertions)]
 42        {
 43            assert_eq!(self.basek(), a.basek());
 44            assert_eq!(a.n(), ks.n());
 45        }
 46        let (mut tmp_glwe, scratch1) = scratch.take_glwe_ct(a.n(), a.basek(), self.k(), 1);
 47        tmp_glwe.keyswitch(module, a, &ks.0, scratch1);
 48        self.sample_extract(&tmp_glwe);
 49    }
 50
 51    pub fn keyswitch<A, DKs, B: Backend>(
 52        &mut self,
 53        module: &Module<B>,
 54        a: &LWECiphertext<A>,
 55        ksk: &LWESwitchingKeyExec<DKs, B>,
 56        scratch: &mut Scratch<B>,
 57    ) where
 58        A: DataRef,
 59        DKs: DataRef,
 60        Module<B>: GLWEKeyswitchFamily<B>,
 61        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
 62    {
 63        #[cfg(debug_assertions)]
 64        {
 65            assert!(self.n() <= module.n());
 66            assert!(a.n() <= module.n());
 67            assert_eq!(self.basek(), a.basek());
 68        }
 69
 70        let max_k: usize = self.k().max(a.k());
 71        let basek: usize = self.basek();
 72
 73        let (mut glwe, scratch1) = scratch.take_glwe_ct(ksk.n(), basek, max_k, 1);
 74        glwe.data.zero();
 75
 76        let n_lwe: usize = a.n();
 77
 78        (0..a.size()).for_each(|i| {
 79            let data_lwe: &[i64] = a.data.at(0, i);
 80            glwe.data.at_mut(0, i)[0] = data_lwe[0];
 81            glwe.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]);
 82        });
 83
 84        glwe.keyswitch_inplace(module, &ksk.0, scratch1);
 85
 86        self.sample_extract(&glwe);
 87    }
 88}
 89
 90impl GLWECiphertext<Vec<u8>> {
 91    pub fn from_lwe_scratch_space<B: Backend>(
 92        module: &Module<B>,
 93        n: usize,
 94        basek: usize,
 95        k_lwe: usize,
 96        k_glwe: usize,
 97        k_ksk: usize,
 98        rank: usize,
 99    ) -> usize
100    where
101        Module<B>: GLWEKeyswitchFamily<B>,
102    {
103        GLWECiphertext::keyswitch_scratch_space(module, n, basek, k_glwe, k_lwe, k_ksk, 1, 1, rank)
104            + GLWECiphertext::bytes_of(n, basek, k_lwe, 1)
105    }
106}
107
108impl<D: DataMut> GLWECiphertext<D> {
109    pub fn from_lwe<DLwe, DKsk, B: Backend>(
110        &mut self,
111        module: &Module<B>,
112        lwe: &LWECiphertext<DLwe>,
113        ksk: &LWEToGLWESwitchingKeyExec<DKsk, B>,
114        scratch: &mut Scratch<B>,
115    ) where
116        DLwe: DataRef,
117        DKsk: DataRef,
118        Module<B>: GLWEKeyswitchFamily<B>,
119        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
120    {
121        #[cfg(debug_assertions)]
122        {
123            assert!(lwe.n() <= self.n());
124            assert_eq!(self.basek(), self.basek());
125        }
126
127        let (mut glwe, scratch1) = scratch.take_glwe_ct(ksk.n(), lwe.basek(), lwe.k(), 1);
128        glwe.data.zero();
129
130        let n_lwe: usize = lwe.n();
131
132        (0..lwe.size()).for_each(|i| {
133            let data_lwe: &[i64] = lwe.data.at(0, i);
134            glwe.data.at_mut(0, i)[0] = data_lwe[0];
135            glwe.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]);
136        });
137
138        self.keyswitch(module, &glwe, &ksk.0, scratch1);
139    }
140}