1use poulpy_hal::{
  2    api::{
  3        ModuleN, ScratchTakeBasic, SvpApplyDftToDftInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigBytesOf,
  4        VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf, VecZnxIdftApplyConsume, VecZnxNormalizeTmpBytes,
  5    },
  6    layouts::{Backend, DataRef, DataViewMut, Module, Scratch},
  7};
  8
  9use crate::layouts::{
 10    GLWE, GLWEInfos, GLWEPlaintextToMut, GLWESecretPrepared, GLWEToRef, LWEInfos, SetGLWEInfos, prepared::GLWESecretPreparedToRef,
 11};
 12
 13impl GLWE<Vec<u8>> {
 14    pub fn decrypt_tmp_bytes<A, M, BE: Backend>(module: &M, a_infos: &A) -> usize
 15    where
 16        A: GLWEInfos,
 17        M: GLWEDecrypt<BE>,
 18    {
 19        module.glwe_decrypt_tmp_bytes(a_infos)
 20    }
 21}
 22
 23impl<DataSelf: DataRef> GLWE<DataSelf> {
 24    pub fn decrypt<P, S, M, BE: Backend>(&self, module: &M, pt: &mut P, sk: &S, scratch: &mut Scratch<BE>)
 25    where
 26        P: GLWEPlaintextToMut + GLWEInfos + SetGLWEInfos,
 27        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
 28        M: GLWEDecrypt<BE>,
 29        Scratch<BE>: ScratchTakeBasic,
 30    {
 31        module.glwe_decrypt(self, pt, sk, scratch);
 32    }
 33}
 34
 35pub trait GLWEDecrypt<BE: Backend> {
 36    fn glwe_decrypt_tmp_bytes<A>(&self, infos: &A) -> usize
 37    where
 38        A: GLWEInfos;
 39
 40    fn glwe_decrypt<R, P, S>(&self, res: &R, pt: &mut P, sk: &S, scratch: &mut Scratch<BE>)
 41    where
 42        R: GLWEToRef + GLWEInfos,
 43        P: GLWEPlaintextToMut + GLWEInfos + SetGLWEInfos,
 44        S: GLWESecretPreparedToRef<BE> + GLWEInfos;
 45}
 46
 47impl<BE: Backend> GLWEDecrypt<BE> for Module<BE>
 48where
 49    Self: ModuleN
 50        + VecZnxDftBytesOf
 51        + VecZnxNormalizeTmpBytes
 52        + VecZnxBigBytesOf
 53        + VecZnxDftApply<BE>
 54        + SvpApplyDftToDftInplace<BE>
 55        + VecZnxIdftApplyConsume<BE>
 56        + VecZnxBigAddInplace<BE>
 57        + VecZnxBigAddSmallInplace<BE>
 58        + VecZnxBigNormalize<BE>,
 59    Scratch<BE>: ScratchTakeBasic,
 60{
 61    fn glwe_decrypt_tmp_bytes<A>(&self, infos: &A) -> usize
 62    where
 63        A: GLWEInfos,
 64    {
 65        let size: usize = infos.size();
 66        (self.vec_znx_normalize_tmp_bytes() | self.bytes_of_vec_znx_dft(1, size)) + self.bytes_of_vec_znx_dft(1, size)
 67    }
 68
 69    fn glwe_decrypt<R, P, S>(&self, res: &R, pt: &mut P, sk: &S, scratch: &mut Scratch<BE>)
 70    where
 71        R: GLWEToRef + GLWEInfos,
 72        P: GLWEPlaintextToMut + GLWEInfos + SetGLWEInfos,
 73        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
 74    {
 75        let res: &GLWE<&[u8]> = &res.to_ref();
 76        let sk: &GLWESecretPrepared<&[u8], BE> = &sk.to_ref();
 77
 78        #[cfg(debug_assertions)]
 79        {
 80            assert_eq!(res.rank(), sk.rank()); //NOTE: res.rank() != res.to_ref().rank() if res is of type GLWETensor
 81            assert_eq!(res.n(), sk.n());
 82            assert_eq!(pt.n(), sk.n());
 83        }
 84
 85        let cols: usize = (res.rank() + 1).into();
 86
 87        let (mut c0_big, scratch_1) = scratch.take_vec_znx_big(self, 1, res.size()); // TODO optimize size when pt << ct
 88        c0_big.data_mut().fill(0);
 89
 90        (1..cols).for_each(|i| {
 91            // ci_dft = DFT(a[i]) * DFT(s[i])
 92            let (mut ci_dft, _) = scratch_1.take_vec_znx_dft(self, 1, res.size()); // TODO optimize size when pt << ct
 93            self.vec_znx_dft_apply(1, 0, &mut ci_dft, 0, res.data(), i);
 94            self.svp_apply_dft_to_dft_inplace(&mut ci_dft, 0, &sk.data, i - 1);
 95            let ci_big = self.vec_znx_idft_apply_consume(ci_dft);
 96
 97            // c0_big += a[i] * s[i]
 98            self.vec_znx_big_add_inplace(&mut c0_big, 0, &ci_big, 0);
 99        });
100
101        // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
102        self.vec_znx_big_add_small_inplace(&mut c0_big, 0, res.data(), 0);
103
104        let pt_base2k: usize = pt.base2k().into();
105
106        // pt = norm(BIG(m + e))
107        self.vec_znx_big_normalize(
108            pt.to_mut().data_mut(),
109            pt_base2k,
110            0,
111            0,
112            &c0_big,
113            res.base2k().into(),
114            0,
115            scratch_1,
116        );
117
118        pt.set_k(pt.k().min(res.k()));
119    }
120}