1use poulpy_hal::{
  2    api::{
  3        ModuleN, ScratchTakeBasic, SvpApplyDftToDft, VecZnxBigBytesOf, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
  4        VecZnxDftApply, VecZnxDftBytesOf, VecZnxIdftApplyTmpA,
  5    },
  6    layouts::{
  7        Backend, Data, DataMut, DataRef, Module, ScalarZnx, ScalarZnxToMut, ScalarZnxToRef, Scratch, ZnxInfos, ZnxView,
  8        ZnxViewMut,
  9    },
 10};
 11
 12use crate::{
 13    GetDistribution, ScratchTakeCore,
 14    dist::Distribution,
 15    layouts::{
 16        Base2K, Degree, GLWEInfos, GLWESecret, GLWESecretPreparedFactory, GLWESecretToMut, GLWESecretToRef, LWEInfos, Rank,
 17        TorusPrecision,
 18    },
 19};
 20
 21pub struct GLWESecretTensor<D: Data> {
 22    pub(crate) data: ScalarZnx<D>,
 23    pub(crate) rank: Rank,
 24    pub(crate) dist: Distribution,
 25}
 26
 27impl GLWESecretTensor<Vec<u8>> {
 28    pub(crate) fn pairs(rank: usize) -> usize {
 29        (((rank + 1) * rank) >> 1).max(1)
 30    }
 31}
 32
 33impl<D: Data> GetDistribution for GLWESecretTensor<D> {
 34    fn dist(&self) -> &Distribution {
 35        &self.dist
 36    }
 37}
 38
 39impl<D: Data> LWEInfos for GLWESecretTensor<D> {
 40    fn base2k(&self) -> Base2K {
 41        Base2K(0)
 42    }
 43
 44    fn k(&self) -> TorusPrecision {
 45        TorusPrecision(0)
 46    }
 47
 48    fn n(&self) -> Degree {
 49        Degree(self.data.n() as u32)
 50    }
 51
 52    fn size(&self) -> usize {
 53        1
 54    }
 55}
 56
 57impl<D: DataRef> GLWESecretTensor<D> {
 58    pub fn at(&self, mut i: usize, mut j: usize) -> ScalarZnx<&[u8]> {
 59        if i > j {
 60            std::mem::swap(&mut i, &mut j);
 61        };
 62        let rank: usize = self.rank().into();
 63        ScalarZnx {
 64            data: bytemuck::cast_slice(self.data.at(i * rank + j - (i * (i + 1) / 2), 0)),
 65            n: self.n().into(),
 66            cols: 1,
 67        }
 68    }
 69}
 70
 71impl<D: DataMut> GLWESecretTensor<D> {
 72    pub fn at_mut(&mut self, mut i: usize, mut j: usize) -> ScalarZnx<&mut [u8]> {
 73        if i > j {
 74            std::mem::swap(&mut i, &mut j);
 75        };
 76        let rank: usize = self.rank().into();
 77        ScalarZnx {
 78            n: self.n().into(),
 79            data: bytemuck::cast_slice_mut(self.data.at_mut(i * rank + j - (i * (i + 1) / 2), 0)),
 80            cols: 1,
 81        }
 82    }
 83}
 84
 85impl<D: Data> GLWEInfos for GLWESecretTensor<D> {
 86    fn rank(&self) -> Rank {
 87        self.rank
 88    }
 89}
 90
 91impl<D: DataRef> GLWESecretToRef for GLWESecretTensor<D> {
 92    fn to_ref(&self) -> GLWESecret<&[u8]> {
 93        GLWESecret {
 94            data: self.data.to_ref(),
 95            dist: self.dist,
 96        }
 97    }
 98}
 99
100impl<D: DataMut> GLWESecretToMut for GLWESecretTensor<D> {
101    fn to_mut(&mut self) -> GLWESecret<&mut [u8]> {
102        GLWESecret {
103            dist: self.dist,
104            data: self.data.to_mut(),
105        }
106    }
107}
108
109impl GLWESecretTensor<Vec<u8>> {
110    pub fn alloc_from_infos<A>(infos: &A) -> Self
111    where
112        A: GLWEInfos,
113    {
114        Self::alloc(infos.n(), infos.rank())
115    }
116
117    pub fn alloc(n: Degree, rank: Rank) -> Self {
118        GLWESecretTensor {
119            data: ScalarZnx::alloc(n.into(), Self::pairs(rank.into())),
120            rank,
121            dist: Distribution::NONE,
122        }
123    }
124
125    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
126    where
127        A: GLWEInfos,
128    {
129        Self::bytes_of(infos.n(), Self::pairs(infos.rank().into()).into())
130    }
131
132    pub fn bytes_of(n: Degree, rank: Rank) -> usize {
133        ScalarZnx::bytes_of(n.into(), Self::pairs(rank.into()))
134    }
135}
136
137impl<D: DataMut> GLWESecretTensor<D> {
138    pub fn prepare<M, S, BE: Backend>(&mut self, module: &M, other: &S, scratch: &mut Scratch<BE>)
139    where
140        M: GLWESecretTensorFactory<BE>,
141        S: GLWESecretToRef + GLWEInfos,
142        Scratch<BE>: ScratchTakeCore<BE>,
143    {
144        module.glwe_secret_tensor_prepare(self, other, scratch);
145    }
146}
147
148pub trait GLWESecretTensorFactory<BE: Backend> {
149    fn glwe_secret_tensor_prepare_tmp_bytes(&self, rank: Rank) -> usize;
150
151    fn glwe_secret_tensor_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<BE>)
152    where
153        R: GLWESecretToMut + GLWEInfos,
154        O: GLWESecretToRef + GLWEInfos;
155}
156
157impl<BE: Backend> GLWESecretTensorFactory<BE> for Module<BE>
158where
159    Self: ModuleN
160        + GLWESecretPreparedFactory<BE>
161        + VecZnxBigNormalize<BE>
162        + VecZnxDftApply<BE>
163        + SvpApplyDftToDft<BE>
164        + VecZnxIdftApplyTmpA<BE>
165        + VecZnxBigNormalize<BE>
166        + VecZnxDftBytesOf
167        + VecZnxBigBytesOf
168        + VecZnxBigNormalizeTmpBytes,
169    Scratch<BE>: ScratchTakeCore<BE>,
170{
171    fn glwe_secret_tensor_prepare_tmp_bytes(&self, rank: Rank) -> usize {
172        self.bytes_of_glwe_secret_prepared(rank)
173            + self.bytes_of_vec_znx_dft(rank.into(), 1)
174            + self.bytes_of_vec_znx_dft(1, 1)
175            + self.bytes_of_vec_znx_big(1, 1)
176            + self.vec_znx_big_normalize_tmp_bytes()
177    }
178
179    fn glwe_secret_tensor_prepare<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
180    where
181        R: GLWESecretToMut + GLWEInfos,
182        A: GLWESecretToRef + GLWEInfos,
183    {
184        let res: &mut GLWESecret<&mut [u8]> = &mut res.to_mut();
185        let a: &GLWESecret<&[u8]> = &a.to_ref();
186
187        assert_eq!(res.rank(), GLWESecretTensor::pairs(a.rank().into()) as u32);
188        assert_eq!(res.n(), self.n() as u32);
189        assert_eq!(a.n(), self.n() as u32);
190
191        let rank: usize = a.rank().into();
192
193        let (mut a_prepared, scratch_1) = scratch.take_glwe_secret_prepared(self, rank.into());
194        a_prepared.prepare(self, a);
195
196        let base2k: usize = 17;
197
198        let (mut a_dft, scratch_2) = scratch_1.take_vec_znx_dft(self, rank, 1);
199        for i in 0..rank {
200            self.vec_znx_dft_apply(1, 0, &mut a_dft, i, &a.data.as_vec_znx(), i);
201        }
202
203        let (mut a_ij_big, scratch_3) = scratch_2.take_vec_znx_big(self, 1, 1);
204        let (mut a_ij_dft, scratch_4) = scratch_3.take_vec_znx_dft(self, 1, 1);
205
206        // sk_tensor = sk (x) sk
207        // For example: (s0, s1) (x) (s0, s1) = (s0^2, s0s1, s1^2)
208        for i in 0..rank {
209            for j in i..rank {
210                let idx: usize = i * rank + j - (i * (i + 1) / 2);
211                self.svp_apply_dft_to_dft(&mut a_ij_dft, 0, &a_prepared.data, j, &a_dft, i);
212                self.vec_znx_idft_apply_tmpa(&mut a_ij_big, 0, &mut a_ij_dft, 0);
213
214                self.vec_znx_big_normalize(
215                    &mut res.data.as_vec_znx_mut(),
216                    base2k,
217                    0,
218                    idx,
219                    &a_ij_big,
220                    base2k,
221                    0,
222                    scratch_4,
223                );
224            }
225        }
226    }
227}