1use poulpy_hal::{
  2    api::{
  3        ScratchAvailable, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftApply,
  4        VecZnxDftBytesOf, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft,
  5        VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
  6    },
  7    layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxView, ZnxViewMut, ZnxZero},
  8};
  9
 10use crate::layouts::{GGLWEInfos, GLWE, GLWELayout, LWE, LWEInfos, Rank, TorusPrecision, prepared::LWESwitchingKeyPrepared};
 11
 12impl LWE<Vec<u8>> {
 13    pub fn keyswitch_scratch_space<B: Backend, OUT, IN, KEY>(
 14        module: &Module<B>,
 15        out_infos: &OUT,
 16        in_infos: &IN,
 17        key_infos: &KEY,
 18    ) -> usize
 19    where
 20        OUT: LWEInfos,
 21        IN: LWEInfos,
 22        KEY: GGLWEInfos,
 23        Module<B>: VecZnxDftBytesOf
 24            + VmpApplyDftToDftTmpBytes
 25            + VecZnxBigNormalizeTmpBytes
 26            + VmpApplyDftToDftTmpBytes
 27            + VmpApplyDftToDft<B>
 28            + VmpApplyDftToDftAdd<B>
 29            + VecZnxDftApply<B>
 30            + VecZnxIdftApplyConsume<B>
 31            + VecZnxBigAddSmallInplace<B>
 32            + VecZnxBigNormalize<B>
 33            + VecZnxNormalizeTmpBytes,
 34    {
 35        let max_k: TorusPrecision = in_infos.k().max(out_infos.k());
 36
 37        let glwe_in_infos: GLWELayout = GLWELayout {
 38            n: module.n().into(),
 39            base2k: in_infos.base2k(),
 40            k: max_k,
 41            rank: Rank(1),
 42        };
 43
 44        let glwe_out_infos: GLWELayout = GLWELayout {
 45            n: module.n().into(),
 46            base2k: out_infos.base2k(),
 47            k: max_k,
 48            rank: Rank(1),
 49        };
 50
 51        let glwe_in: usize = GLWE::bytes_of_from_infos(module, &glwe_in_infos);
 52        let glwe_out: usize = GLWE::bytes_of_from_infos(module, &glwe_out_infos);
 53        let ks: usize = GLWE::keyswitch_scratch_space(module, &glwe_out_infos, &glwe_in_infos, key_infos);
 54
 55        glwe_in + glwe_out + ks
 56    }
 57}
 58
 59impl<DLwe: DataMut> LWE<DLwe> {
 60    pub fn keyswitch<A, DKs, B: Backend>(
 61        &mut self,
 62        module: &Module<B>,
 63        a: &LWE<A>,
 64        ksk: &LWESwitchingKeyPrepared<DKs, B>,
 65        scratch: &mut Scratch<B>,
 66    ) where
 67        A: DataRef,
 68        DKs: DataRef,
 69        Module<B>: VecZnxDftBytesOf
 70            + VmpApplyDftToDftTmpBytes
 71            + VecZnxBigNormalizeTmpBytes
 72            + VmpApplyDftToDft<B>
 73            + VmpApplyDftToDftAdd<B>
 74            + VecZnxDftApply<B>
 75            + VecZnxIdftApplyConsume<B>
 76            + VecZnxBigAddSmallInplace<B>
 77            + VecZnxBigNormalize<B>
 78            + VecZnxNormalize<B>
 79            + VecZnxNormalizeTmpBytes
 80            + VecZnxCopy,
 81        Scratch<B>: ScratchAvailable,
 82    {
 83        #[cfg(debug_assertions)]
 84        {
 85            assert!(self.n() <= module.n() as u32);
 86            assert!(a.n() <= module.n() as u32);
 87            assert!(scratch.available() >= LWE::keyswitch_scratch_space(module, self, a, ksk));
 88        }
 89
 90        let max_k: TorusPrecision = self.k().max(a.k());
 91
 92        let a_size: usize = a.k().div_ceil(ksk.base2k()) as usize;
 93
 94        let (mut glwe_in, scratch_1) = scratch.take_glwe_ct(&GLWELayout {
 95            n: ksk.n(),
 96            base2k: a.base2k(),
 97            k: max_k,
 98            rank: Rank(1),
 99        });
100        glwe_in.data.zero();
101
102        let (mut glwe_out, scratch_1) = scratch_1.take_glwe_ct(&GLWELayout {
103            n: ksk.n(),
104            base2k: self.base2k(),
105            k: max_k,
106            rank: Rank(1),
107        });
108
109        let n_lwe: usize = a.n().into();
110
111        for i in 0..a_size {
112            let data_lwe: &[i64] = a.data.at(0, i);
113            glwe_in.data.at_mut(0, i)[0] = data_lwe[0];
114            glwe_in.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]);
115        }
116
117        glwe_out.keyswitch(module, &glwe_in, &ksk.0, scratch_1);
118        self.sample_extract(&glwe_out);
119    }
120}