1use poulpy_hal::{
 2    api::{VecZnxNormalize, VecZnxNormalizeTmpBytes},
 3    layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxView, ZnxViewMut},
 4};
 5
 6use crate::{
 7    ScratchTakeCore,
 8    layouts::{LWE, LWEInfos, LWEPlaintext, LWEPlaintextToMut, LWESecret, LWESecretToRef, LWEToRef, SetLWEInfos, TorusPrecision},
 9};
10
11impl<DataSelf: DataRef + DataMut> LWE<DataSelf> {
12    pub fn decrypt<P, S, M, BE: Backend>(&self, module: &M, pt: &mut P, sk: &S, scratch: &mut Scratch<BE>)
13    where
14        P: LWEPlaintextToMut + SetLWEInfos + LWEInfos,
15        S: LWESecretToRef,
16        M: LWEDecrypt<BE>,
17        Scratch<BE>: ScratchTakeCore<BE>,
18    {
19        module.lwe_decrypt(self, pt, sk, scratch);
20    }
21}
22
23pub trait LWEDecrypt<BE: Backend> {
24    fn lwe_decrypt<R, P, S>(&self, res: &R, pt: &mut P, sk: &S, scratch: &mut Scratch<BE>)
25    where
26        R: LWEToRef,
27        P: LWEPlaintextToMut + SetLWEInfos + LWEInfos,
28        S: LWESecretToRef,
29        Scratch<BE>: ScratchTakeCore<BE>;
30    fn lwe_decrypt_tmp_bytes<A>(&self, infos: &A) -> usize
31    where
32        A: LWEInfos;
33}
34
35impl<BE: Backend> LWEDecrypt<BE> for Module<BE>
36where
37    Self: Sized + VecZnxNormalize<BE> + VecZnxNormalizeTmpBytes,
38{
39    fn lwe_decrypt_tmp_bytes<A>(&self, infos: &A) -> usize
40    where
41        A: LWEInfos,
42    {
43        self.vec_znx_normalize_tmp_bytes() + LWEPlaintext::bytes_of(infos.size())
44    }
45
46    fn lwe_decrypt<R, P, S>(&self, res: &R, pt: &mut P, sk: &S, scratch: &mut Scratch<BE>)
47    where
48        R: LWEToRef,
49        P: LWEPlaintextToMut + SetLWEInfos + LWEInfos,
50        S: LWESecretToRef,
51        Scratch<BE>: ScratchTakeCore<BE>,
52    {
53        let res: &LWE<&[u8]> = &res.to_ref();
54        let sk: LWESecret<&[u8]> = sk.to_ref();
55
56        #[cfg(debug_assertions)]
57        {
58            assert_eq!(res.n(), sk.n());
59        }
60
61        let (mut tmp, scratch_1) = scratch.take_lwe_plaintext(res);
62        for i in 0..res.size() {
63            tmp.data.at_mut(0, i)[0] = res.data.at(0, i)[0]
64                + res.data.at(0, i)[1..]
65                    .iter()
66                    .zip(sk.data.at(0, 0))
67                    .map(|(x, y)| x * y)
68                    .sum::<i64>();
69        }
70
71        let pt_base2k = pt.base2k().into();
72        let res_base2k = res.base2k().into();
73        self.vec_znx_normalize(&mut pt.to_mut().data, pt_base2k, 0, 0, tmp.data(), res_base2k, 0, scratch_1);
74
75        pt.set_k(TorusPrecision(res.k().0.min(pt.size() as u32 * res.base2k().0)));
76    }
77}