1use poulpy_hal::{
  2    api::{
  3        VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftApply, VecZnxDftBytesOf,
  4        VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd,
  5        VmpApplyDftToDftTmpBytes,
  6    },
  7    layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxView, ZnxViewMut, ZnxZero},
  8};
  9
 10use crate::layouts::{GGLWEInfos, GLWE, GLWEInfos, GLWELayout, LWE, LWEInfos, Rank, prepared::GLWEToLWESwitchingKeyPrepared};
 11
 12impl LWE<Vec<u8>> {
 13    pub fn from_glwe_scratch_space<B: Backend, OUT, IN, KEY>(
 14        module: &Module<B>,
 15        lwe_infos: &OUT,
 16        glwe_infos: &IN,
 17        key_infos: &KEY,
 18    ) -> usize
 19    where
 20        OUT: LWEInfos,
 21        IN: GLWEInfos,
 22        KEY: GGLWEInfos,
 23        Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
 24    {
 25        let glwe_layout: GLWELayout = GLWELayout {
 26            n: module.n().into(),
 27            base2k: lwe_infos.base2k(),
 28            k: lwe_infos.k(),
 29            rank: Rank(1),
 30        };
 31
 32        GLWE::bytes_of(
 33            module.n().into(),
 34            lwe_infos.base2k(),
 35            lwe_infos.k(),
 36            1u32.into(),
 37        ) + GLWE::keyswitch_scratch_space(module, &glwe_layout, glwe_infos, key_infos)
 38    }
 39}
 40
 41impl<DLwe: DataMut> LWE<DLwe> {
 42    pub fn sample_extract<DGlwe: DataRef>(&mut self, a: &GLWE<DGlwe>) {
 43        #[cfg(debug_assertions)]
 44        {
 45            assert!(self.n() <= a.n());
 46            assert!(self.base2k() == a.base2k());
 47        }
 48
 49        let min_size: usize = self.size().min(a.size());
 50        let n: usize = self.n().into();
 51
 52        self.data.zero();
 53        (0..min_size).for_each(|i| {
 54            let data_lwe: &mut [i64] = self.data.at_mut(0, i);
 55            data_lwe[0] = a.data.at(0, i)[0];
 56            data_lwe[1..].copy_from_slice(&a.data.at(1, i)[..n]);
 57        });
 58    }
 59
 60    pub fn from_glwe<DGlwe, DKs, B: Backend>(
 61        &mut self,
 62        module: &Module<B>,
 63        a: &GLWE<DGlwe>,
 64        ks: &GLWEToLWESwitchingKeyPrepared<DKs, B>,
 65        scratch: &mut Scratch<B>,
 66    ) where
 67        DGlwe: 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        Scratch<B>:,
 81    {
 82        #[cfg(debug_assertions)]
 83        {
 84            assert_eq!(a.n(), module.n() as u32);
 85            assert_eq!(ks.n(), module.n() as u32);
 86            assert!(self.n() <= module.n() as u32);
 87        }
 88
 89        let glwe_layout: GLWELayout = GLWELayout {
 90            n: module.n().into(),
 91            base2k: self.base2k(),
 92            k: self.k(),
 93            rank: Rank(1),
 94        };
 95
 96        let (mut tmp_glwe, scratch_1) = scratch.take_glwe_ct(&glwe_layout);
 97        tmp_glwe.keyswitch(module, a, &ks.0, scratch_1);
 98        self.sample_extract(&tmp_glwe);
 99    }
100}