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