1use poulpy_hal::{
 2    api::VecZnxNormalizeInplace,
 3    layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxView, ZnxViewMut},
 4};
 5
 6use crate::{
 7    ScratchTakeCore,
 8    layouts::{LWE, LWEInfos, LWEPlaintext, LWEPlaintextToMut, LWESecret, LWESecretToRef, LWEToMut},
 9};
10
11impl<DataSelf: DataRef + DataMut> LWE<DataSelf> {
12    pub fn decrypt<P, S, M, BE: Backend>(&mut self, module: &M, pt: &mut P, sk: &S, scratch: &mut Scratch<BE>)
13    where
14        P: LWEPlaintextToMut,
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: &mut R, pt: &mut P, sk: &S, scratch: &mut Scratch<BE>)
25    where
26        R: LWEToMut,
27        P: LWEPlaintextToMut,
28        S: LWESecretToRef,
29        Scratch<BE>: ScratchTakeCore<BE>;
30}
31
32impl<BE: Backend> LWEDecrypt<BE> for Module<BE>
33where
34    Self: Sized + VecZnxNormalizeInplace<BE>,
35{
36    fn lwe_decrypt<R, P, S>(&self, res: &mut R, pt: &mut P, sk: &S, scratch: &mut Scratch<BE>)
37    where
38        R: LWEToMut,
39        P: LWEPlaintextToMut,
40        S: LWESecretToRef,
41        Scratch<BE>: ScratchTakeCore<BE>,
42    {
43        let res: &mut LWE<&mut [u8]> = &mut res.to_mut();
44        let pt: &mut LWEPlaintext<&mut [u8]> = &mut pt.to_mut();
45        let sk: LWESecret<&[u8]> = sk.to_ref();
46
47        #[cfg(debug_assertions)]
48        {
49            assert_eq!(res.n(), sk.n());
50        }
51
52        (0..pt.size().min(res.size())).for_each(|i| {
53            pt.data.at_mut(0, i)[0] = res.data.at(0, i)[0]
54                + res.data.at(0, i)[1..]
55                    .iter()
56                    .zip(sk.data.at(0, 0))
57                    .map(|(x, y)| x * y)
58                    .sum::<i64>();
59        });
60        self.vec_znx_normalize_inplace(res.base2k().into(), &mut pt.data, 0, scratch);
61        pt.base2k = res.base2k();
62        pt.k = crate::layouts::TorusPrecision(res.k().0.min(pt.size() as u32 * res.base2k().0));
63    }
64}