1use backend::hal::{
 2    api::{
 3        DataViewMut, SvpApplyInplace, TakeVecZnxBig, TakeVecZnxDft, VecZnxBigAddInplace, VecZnxBigAddSmallInplace,
 4        VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume,
 5        VecZnxNormalizeTmpBytes,
 6    },
 7    layouts::{Backend, DataMut, DataRef, Module, Scratch},
 8};
 9
10use crate::{GLWECiphertext, GLWEPlaintext, GLWESecretExec, Infos};
11
12pub trait GLWEDecryptFamily<B: Backend> = VecZnxDftAllocBytes
13    + VecZnxBigAllocBytes
14    + VecZnxDftFromVecZnx<B>
15    + SvpApplyInplace<B>
16    + VecZnxDftToVecZnxBigConsume<B>
17    + VecZnxBigAddInplace<B>
18    + VecZnxBigAddSmallInplace<B>
19    + VecZnxBigNormalize<B>
20    + VecZnxNormalizeTmpBytes;
21
22impl GLWECiphertext<Vec<u8>> {
23    pub fn decrypt_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize) -> usize
24    where
25        Module<B>: GLWEDecryptFamily<B>,
26    {
27        let size: usize = k.div_ceil(basek);
28        (module.vec_znx_normalize_tmp_bytes(module.n()) | module.vec_znx_dft_alloc_bytes(1, size))
29            + module.vec_znx_dft_alloc_bytes(1, size)
30    }
31}
32
33impl<DataSelf: DataRef> GLWECiphertext<DataSelf> {
34    pub fn decrypt<DataPt: DataMut, DataSk: DataRef, B: Backend>(
35        &self,
36        module: &Module<B>,
37        pt: &mut GLWEPlaintext<DataPt>,
38        sk: &GLWESecretExec<DataSk, B>,
39        scratch: &mut Scratch<B>,
40    ) where
41        Module<B>: GLWEDecryptFamily<B>,
42        Scratch<B>: TakeVecZnxDft<B> + TakeVecZnxBig<B>,
43    {
44        #[cfg(debug_assertions)]
45        {
46            assert_eq!(self.rank(), sk.rank());
47            assert_eq!(self.n(), module.n());
48            assert_eq!(pt.n(), module.n());
49            assert_eq!(sk.n(), module.n());
50        }
51
52        let cols: usize = self.rank() + 1;
53
54        let (mut c0_big, scratch_1) = scratch.take_vec_znx_big(module, 1, self.size()); // TODO optimize size when pt << ct
55        c0_big.data_mut().fill(0);
56
57        {
58            (1..cols).for_each(|i| {
59                // ci_dft = DFT(a[i]) * DFT(s[i])
60                let (mut ci_dft, _) = scratch_1.take_vec_znx_dft(module, 1, self.size()); // TODO optimize size when pt << ct
61                module.vec_znx_dft_from_vec_znx(1, 0, &mut ci_dft, 0, &self.data, i);
62                module.svp_apply_inplace(&mut ci_dft, 0, &sk.data, i - 1);
63                let ci_big = module.vec_znx_dft_to_vec_znx_big_consume(ci_dft);
64
65                // c0_big += a[i] * s[i]
66                module.vec_znx_big_add_inplace(&mut c0_big, 0, &ci_big, 0);
67            });
68        }
69
70        // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
71        module.vec_znx_big_add_small_inplace(&mut c0_big, 0, &self.data, 0);
72
73        // pt = norm(BIG(m + e))
74        module.vec_znx_big_normalize(self.basek(), &mut pt.data, 0, &mut c0_big, 0, scratch_1);
75
76        pt.basek = self.basek();
77        pt.k = pt.k().min(self.k());
78    }
79}