1use poulpy_hal::{
 2    api::{ScratchTakeBasic, SvpPPolBytesOf},
 3    layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxView, ZnxViewMut},
 4};
 5
 6use crate::{
 7    GLWEDecrypt, ScratchTakeCore,
 8    layouts::{GLWEInfos, GLWEPlaintext, GLWESecretPrepared, GLWESecretTensor, GLWESecretTensorPrepared, GLWETensor},
 9};
10
11impl GLWETensor<Vec<u8>> {
12    pub fn decrypt_tmp_bytes<A, M, BE: Backend>(module: &M, a_infos: &A) -> usize
13    where
14        A: GLWEInfos,
15        M: GLWETensorDecrypt<BE>,
16    {
17        module.glwe_tensor_decrypt_tmp_bytes(a_infos)
18    }
19}
20
21impl<DataSelf: DataRef> GLWETensor<DataSelf> {
22    pub fn decrypt<P, S0, S1, M, BE: Backend>(
23        &self,
24        module: &M,
25        pt: &mut GLWEPlaintext<P>,
26        sk: &GLWESecretPrepared<S0, BE>,
27        sk_tensor: &GLWESecretTensorPrepared<S1, BE>,
28        scratch: &mut Scratch<BE>,
29    ) where
30        P: DataMut,
31        S0: DataRef,
32        S1: DataRef,
33        M: GLWETensorDecrypt<BE>,
34        Scratch<BE>: ScratchTakeBasic,
35    {
36        module.glwe_tensor_decrypt(self, pt, sk, sk_tensor, scratch);
37    }
38}
39
40pub trait GLWETensorDecrypt<BE: Backend> {
41    fn glwe_tensor_decrypt_tmp_bytes<A>(&self, infos: &A) -> usize
42    where
43        A: GLWEInfos;
44    fn glwe_tensor_decrypt<R, P, S0, S1>(
45        &self,
46        res: &GLWETensor<R>,
47        pt: &mut GLWEPlaintext<P>,
48        sk: &GLWESecretPrepared<S0, BE>,
49        sk_tensor: &GLWESecretTensorPrepared<S1, BE>,
50        scratch: &mut Scratch<BE>,
51    ) where
52        R: DataRef,
53        P: DataMut,
54        S0: DataRef,
55        S1: DataRef;
56}
57
58impl<BE: Backend> GLWETensorDecrypt<BE> for Module<BE>
59where
60    Self: GLWEDecrypt<BE> + SvpPPolBytesOf,
61    Scratch<BE>: ScratchTakeCore<BE>,
62{
63    fn glwe_tensor_decrypt_tmp_bytes<A>(&self, infos: &A) -> usize
64    where
65        A: GLWEInfos,
66    {
67        self.glwe_decrypt_tmp_bytes(infos)
68    }
69
70    fn glwe_tensor_decrypt<R, P, S0, S1>(
71        &self,
72        res: &GLWETensor<R>,
73        pt: &mut GLWEPlaintext<P>,
74        sk: &GLWESecretPrepared<S0, BE>,
75        sk_tensor: &GLWESecretTensorPrepared<S1, BE>,
76        scratch: &mut Scratch<BE>,
77    ) where
78        R: DataRef,
79        P: DataMut,
80        S0: DataRef,
81        S1: DataRef,
82    {
83        let rank: usize = sk.rank().as_usize();
84
85        let (mut sk_grouped, scratch_1) = scratch.take_glwe_secret_prepared(self, (GLWESecretTensor::pairs(rank) + rank).into());
86
87        for i in 0..rank {
88            sk_grouped.data.at_mut(i, 0).copy_from_slice(sk.data.at(i, 0));
89        }
90
91        for i in 0..sk_grouped.rank().as_usize() - rank {
92            sk_grouped.data.at_mut(i + rank, 0).copy_from_slice(sk_tensor.data.at(i, 0));
93        }
94
95        self.glwe_decrypt(res, pt, &sk_grouped, scratch_1);
96    }
97}