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