1use poulpy_hal::{
 2    api::{ScratchOwnedAlloc, ScratchOwnedBorrow, ZnNormalizeInplace},
 3    layouts::{Backend, DataMut, DataRef, Module, ScratchOwned, ZnxView, ZnxViewMut},
 4    oep::{ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl},
 5};
 6
 7use crate::layouts::{LWE, LWEInfos, LWEPlaintext, LWEPlaintextToMut, LWESecret, LWESecretToRef, LWEToMut};
 8
 9impl<DataSelf: DataRef + DataMut> LWE<DataSelf> {
10    pub fn decrypt<P, S, M, B>(&mut self, module: &M, pt: &mut P, sk: S)
11    where
12        P: LWEPlaintextToMut,
13        S: LWESecretToRef,
14        M: LWEDecrypt<B>,
15        B: Backend + ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
16    {
17        module.lwe_decrypt(self, pt, sk);
18    }
19}
20
21pub trait LWEDecrypt<BE: Backend>
22where
23    Self: Sized + ZnNormalizeInplace<BE>,
24{
25    fn lwe_decrypt<R, P, S>(&self, res: &mut R, pt: &mut P, sk: S)
26    where
27        R: LWEToMut,
28        P: LWEPlaintextToMut,
29        S: LWESecretToRef,
30        BE: Backend + ScratchOwnedAllocImpl<BE> + ScratchOwnedBorrowImpl<BE>,
31    {
32        let res: &mut LWE<&mut [u8]> = &mut res.to_mut();
33        let pt: &mut LWEPlaintext<&mut [u8]> = &mut pt.to_mut();
34        let sk: LWESecret<&[u8]> = sk.to_ref();
35
36        #[cfg(debug_assertions)]
37        {
38            assert_eq!(res.n(), sk.n());
39        }
40
41        (0..pt.size().min(res.size())).for_each(|i| {
42            pt.data.at_mut(0, i)[0] = res.data.at(0, i)[0]
43                + res.data.at(0, i)[1..]
44                    .iter()
45                    .zip(sk.data.at(0, 0))
46                    .map(|(x, y)| x * y)
47                    .sum::<i64>();
48        });
49        self.zn_normalize_inplace(
50            1,
51            res.base2k().into(),
52            &mut pt.data,
53            0,
54            ScratchOwned::alloc(size_of::<i64>()).borrow(),
55        );
56        pt.base2k = res.base2k();
57        pt.k = crate::layouts::TorusPrecision(res.k().0.min(pt.size() as u32 * res.base2k().0));
58    }
59}
60
61impl<BE: Backend> LWEDecrypt<BE> for Module<BE> where Self: Sized + ZnNormalizeInplace<BE> {}