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