1use poulpy_hal::{
  2    api::{VecZnxDftAlloc, VecZnxDftApply, VecZnxDftBytesOf},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, VecZnxDft, VecZnxDftToMut, VecZnxDftToRef, ZnxInfos},
  4};
  5
  6use crate::layouts::{Base2K, Degree, GLWE, GLWEInfos, GLWEToRef, GetDegree, LWEInfos, Rank, TorusPrecision};
  7
  8#[derive(PartialEq, Eq)]
  9pub struct GLWEPrepared<D: Data, B: Backend> {
 10    pub(crate) data: VecZnxDft<D, B>,
 11    pub(crate) base2k: Base2K,
 12    pub(crate) k: TorusPrecision,
 13}
 14
 15impl<D: Data, B: Backend> LWEInfos for GLWEPrepared<D, B> {
 16    fn base2k(&self) -> Base2K {
 17        self.base2k
 18    }
 19
 20    fn k(&self) -> TorusPrecision {
 21        self.k
 22    }
 23
 24    fn size(&self) -> usize {
 25        self.data.size()
 26    }
 27
 28    fn n(&self) -> Degree {
 29        Degree(self.data.n() as u32)
 30    }
 31}
 32
 33impl<D: Data, B: Backend> GLWEInfos for GLWEPrepared<D, B> {
 34    fn rank(&self) -> Rank {
 35        Rank(self.data.cols() as u32 - 1)
 36    }
 37}
 38
 39pub trait GLWEPreparedAlloc<B: Backend>
 40where
 41    Self: GetDegree + VecZnxDftAlloc<B> + VecZnxDftBytesOf + VecZnxDftApply<B>,
 42{
 43    fn alloc_glwe_prepared(&self, base2k: Base2K, k: TorusPrecision, rank: Rank) -> GLWEPrepared<Vec<u8>, B> {
 44        GLWEPrepared {
 45            data: self.vec_znx_dft_alloc((rank + 1).into(), k.0.div_ceil(base2k.0) as usize),
 46            base2k,
 47            k,
 48        }
 49    }
 50
 51    fn alloc_glwe_prepared_from_infos<A>(&self, infos: &A) -> GLWEPrepared<Vec<u8>, B>
 52    where
 53        A: GLWEInfos,
 54    {
 55        self.alloc_glwe_prepared(infos.base2k(), infos.k(), infos.rank())
 56    }
 57
 58    fn bytes_of_glwe_prepared(&self, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize {
 59        self.bytes_of_vec_znx_dft((rank + 1).into(), k.0.div_ceil(base2k.0) as usize)
 60    }
 61
 62    fn bytes_of_glwe_prepared_from_infos<A>(&self, infos: &A) -> usize
 63    where
 64        A: GLWEInfos,
 65    {
 66        self.bytes_of_glwe_prepared(infos.base2k(), infos.k(), infos.rank())
 67    }
 68
 69    fn prepare_glwe<R, O>(&self, res: &mut R, other: &O)
 70    where
 71        R: GLWEPreparedToMut<B>,
 72        O: GLWEToRef,
 73    {
 74        {
 75            let mut res: GLWEPrepared<&mut [u8], B> = res.to_mut();
 76            let other: GLWE<&[u8]> = other.to_ref();
 77
 78            assert_eq!(res.n(), self.ring_degree());
 79            assert_eq!(other.n(), self.ring_degree());
 80            assert_eq!(res.size(), other.size());
 81            assert_eq!(res.k(), other.k());
 82            assert_eq!(res.base2k(), other.base2k());
 83
 84            for i in 0..(res.rank() + 1).into() {
 85                self.vec_znx_dft_apply(1, 0, &mut res.data, i, &other.data, i);
 86            }
 87        }
 88    }
 89}
 90
 91impl<B: Backend> GLWEPreparedAlloc<B> for Module<B> where Self: VecZnxDftAlloc<B> + VecZnxDftBytesOf + VecZnxDftApply<B> {}
 92
 93impl<B: Backend> GLWEPrepared<Vec<u8>, B> {
 94    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
 95    where
 96        A: GLWEInfos,
 97        M: GLWEPreparedAlloc<B>,
 98    {
 99        module.alloc_glwe_prepared_from_infos(infos)
100    }
101
102    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self
103    where
104        M: GLWEPreparedAlloc<B>,
105    {
106        module.alloc_glwe_prepared(base2k, k, rank)
107    }
108
109    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
110    where
111        A: GLWEInfos,
112        M: GLWEPreparedAlloc<B>,
113    {
114        module.bytes_of_glwe_prepared_from_infos(infos)
115    }
116
117    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize
118    where
119        M: GLWEPreparedAlloc<B>,
120    {
121        module.bytes_of_glwe_prepared(base2k, k, rank)
122    }
123}
124
125impl<D: DataMut, B: Backend> GLWEPrepared<D, B> {
126    pub fn prepare<O, M>(&mut self, module: &M, other: &O)
127    where
128        O: GLWEToRef,
129        M: GLWEPreparedAlloc<B>,
130    {
131        module.prepare_glwe(self, other);
132    }
133}
134
135pub trait GLWEPreparedToMut<B: Backend> {
136    fn to_mut(&mut self) -> GLWEPrepared<&mut [u8], B>;
137}
138
139impl<D: DataMut, B: Backend> GLWEPreparedToMut<B> for GLWEPrepared<D, B> {
140    fn to_mut(&mut self) -> GLWEPrepared<&mut [u8], B> {
141        GLWEPrepared {
142            k: self.k,
143            base2k: self.base2k,
144            data: self.data.to_mut(),
145        }
146    }
147}
148
149pub trait GLWEPreparedToRef<B: Backend> {
150    fn to_ref(&self) -> GLWEPrepared<&[u8], B>;
151}
152
153impl<D: DataRef, B: Backend> GLWEPreparedToRef<B> for GLWEPrepared<D, B> {
154    fn to_ref(&self) -> GLWEPrepared<&[u8], B> {
155        GLWEPrepared {
156            data: self.data.to_ref(),
157            k: self.k,
158            base2k: self.base2k,
159        }
160    }
161}