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