1use poulpy_hal::{
  2    api::{SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, SvpPPol, SvpPPolToMut, SvpPPolToRef, ZnxInfos},
  4};
  5
  6use crate::{
  7    dist::Distribution,
  8    layouts::{
  9        Base2K, Degree, GLWEInfos, GLWESecret, GLWESecretToMut, GLWESecretToRef, LWEInfos, Rank, TorusPrecision,
 10        prepared::SetDist,
 11    },
 12};
 13
 14pub struct GLWESecretPrepared<D: Data, B: Backend> {
 15    pub(crate) data: SvpPPol<D, B>,
 16    pub(crate) dist: Distribution,
 17}
 18
 19impl<D: DataRef, B: Backend> SetDist for GLWESecretPrepared<D, B> {
 20    fn set_dist(&mut self, dist: Distribution) {
 21        self.dist = dist
 22    }
 23}
 24
 25impl<D: Data, B: Backend> LWEInfos for GLWESecretPrepared<D, B> {
 26    fn base2k(&self) -> Base2K {
 27        Base2K(0)
 28    }
 29
 30    fn k(&self) -> TorusPrecision {
 31        TorusPrecision(0)
 32    }
 33
 34    fn n(&self) -> Degree {
 35        Degree(self.data.n() as u32)
 36    }
 37
 38    fn size(&self) -> usize {
 39        self.data.size()
 40    }
 41}
 42impl<D: Data, B: Backend> GLWEInfos for GLWESecretPrepared<D, B> {
 43    fn rank(&self) -> Rank {
 44        Rank(self.data.cols() as u32)
 45    }
 46}
 47impl<B: Backend> GLWESecretPrepared<Vec<u8>, B> {
 48    pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
 49    where
 50        A: GLWEInfos,
 51        Module<B>: SvpPPolAlloc<B>,
 52    {
 53        assert_eq!(module.n() as u32, infos.n());
 54        Self::alloc_with(module, infos.rank())
 55    }
 56
 57    pub fn alloc_with(module: &Module<B>, rank: Rank) -> Self
 58    where
 59        Module<B>: SvpPPolAlloc<B>,
 60    {
 61        Self {
 62            data: module.svp_ppol_alloc(rank.into()),
 63            dist: Distribution::NONE,
 64        }
 65    }
 66
 67    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
 68    where
 69        A: GLWEInfos,
 70        Module<B>: SvpPPolAllocBytes,
 71    {
 72        assert_eq!(module.n() as u32, infos.n());
 73        Self::alloc_bytes_with(module, infos.rank())
 74    }
 75
 76    pub fn alloc_bytes_with(module: &Module<B>, rank: Rank) -> usize
 77    where
 78        Module<B>: SvpPPolAllocBytes,
 79    {
 80        module.svp_ppol_alloc_bytes(rank.into())
 81    }
 82}
 83
 84impl<D: Data, B: Backend> GLWESecretPrepared<D, B> {
 85    pub fn n(&self) -> Degree {
 86        Degree(self.data.n() as u32)
 87    }
 88
 89    pub fn rank(&self) -> Rank {
 90        Rank(self.data.cols() as u32)
 91    }
 92}
 93
 94pub trait GLWESecretPrepareTmpBytes {
 95    fn glwe_secret_prepare_tmp_bytes<A>(&self, infos: &A)
 96    where
 97        A: GLWEInfos;
 98}
 99
100impl<B: Backend> GLWESecretPrepareTmpBytes for Module<B> {
101    fn glwe_secret_prepare_tmp_bytes<A>(&self, infos: &A)
102    where
103        A: GLWEInfos,
104    {
105        0
106    }
107}
108
109impl<B: Backend> GLWESecretPrepared<Vec<u8>, B>
110where
111    Module<B>: GLWESecretPrepareTmpBytes,
112{
113    fn prepare_tmp_bytes<A>(&self, module: &Module<B>, infos: &A) -> usize
114    where
115        A: GLWEInfos,
116    {
117        0
118    }
119}
120
121pub trait GLWESecretPrepare<B: Backend> {
122    fn glwe_secret_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
123    where
124        R: GLWESecretPreparedToMut<B> + SetDist,
125        O: GLWESecretToRef;
126}
127
128impl<B: Backend> GLWESecretPrepare<B> for Module<B>
129where
130    Module<B>: SvpPrepare<B>,
131{
132    fn glwe_secret_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
133    where
134        R: GLWESecretPreparedToMut<B> + SetDist,
135        O: GLWESecretToRef,
136    {
137        {
138            let res: GLWESecretPrepared<&mut [u8], _> = res.to_mut();
139            let other: GLWESecret<&[u8]> = other.to_ref();
140
141            for i in 0..self.rank().into() {
142                self.svp_prepare(&mut res.data, i, &other.data, i);
143            }
144        }
145
146        res.set_dist(other.dist);
147    }
148}
149
150pub trait GLWESecretPrepareAlloc<B: Backend> {
151    fn glwe_secret_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
152    where
153        O: GLWESecretToMut;
154}
155
156impl<B: Backend> GLWESecretPrepareAlloc<B> for Module<B>
157where
158    Module<B>: GLWESecretPrepare<B>,
159{
160    fn glwe_secret_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
161    where
162        O: GLWESecretToMut,
163    {
164        let mut ct_prep: GLWESecretPrepared<Vec<u8>, B> = GLWESecretPrepared::alloc(self, self);
165        self.glwe_secret_prepare(&mut ct_prep, other, scratch);
166        ct_prep
167    }
168}
169
170pub trait GLWESecretPreparedToRef<B: Backend> {
171    fn to_ref(&self) -> GLWESecretPrepared<&[u8], B>;
172}
173
174impl<D: DataRef, B: Backend> GLWESecretPreparedToRef<B> for GLWESecretPrepared<D, B> {
175    fn to_ref(&self) -> GLWESecretPrepared<&[u8], B> {
176        GLWESecretPrepared {
177            data: self.data.to_ref(),
178            dist: self.dist,
179        }
180    }
181}
182
183pub trait GLWESecretPreparedToMut<B: Backend> {
184    fn to_mut(&mut self) -> GLWESecretPrepared<&mut [u8], B>;
185}
186
187impl<D: DataMut, B: Backend> GLWESecretPreparedToMut<B> for GLWESecretPrepared<D, B> {
188    fn to_mut(&mut self) -> GLWESecretPrepared<&mut [u8], B> {
189        GLWESecretPrepared {
190            dist: self.dist,
191            data: self.data.to_mut(),
192        }
193    }
194}