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