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