arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/core/src/glwe/ops.rs
core/src/glwe/ops.rsRaw
  1use backend::{FFT64, Module, Scratch, VecZnx, VecZnxOps, ZnxZero};
  2
  3use crate::{GLWECiphertext, GLWECiphertextToMut, GLWECiphertextToRef, Infos, SetMetaData};
  4
  5pub trait GLWEOps: GLWECiphertextToMut + SetMetaData {
  6    fn add<A, B>(&mut self, module: &Module<FFT64>, a: &A, b: &B)
  7    where
  8        A: GLWECiphertextToRef,
  9        B: GLWECiphertextToRef,
 10    {
 11        #[cfg(debug_assertions)]
 12        {
 13            assert_eq!(a.n(), module.n());
 14            assert_eq!(b.n(), module.n());
 15            assert_eq!(self.n(), module.n());
 16            assert_eq!(a.basek(), b.basek());
 17            assert_eq!(self.basek(), a.basek());
 18            assert!(self.rank() >= a.rank().max(b.rank()));
 19        }
 20
 21        let min_col: usize = a.rank().min(b.rank()) + 1;
 22        let max_col: usize = a.rank().max(b.rank() + 1);
 23        let self_col: usize = self.rank() + 1;
 24
 25        let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
 26        let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
 27        let b_ref: &GLWECiphertext<&[u8]> = &b.to_ref();
 28
 29        (0..min_col).for_each(|i| {
 30            module.vec_znx_add(&mut self_mut.data, i, &a_ref.data, i, &b_ref.data, i);
 31        });
 32
 33        if a.rank() > b.rank() {
 34            (min_col..max_col).for_each(|i| {
 35                module.vec_znx_copy(&mut self_mut.data, i, &a_ref.data, i);
 36            });
 37        } else {
 38            (min_col..max_col).for_each(|i| {
 39                module.vec_znx_copy(&mut self_mut.data, i, &b_ref.data, i);
 40            });
 41        }
 42
 43        let size: usize = self_mut.size();
 44        (max_col..self_col).for_each(|i| {
 45            (0..size).for_each(|j| {
 46                self_mut.data.zero_at(i, j);
 47            });
 48        });
 49
 50        // self.set_basek(a.basek());
 51        // self.set_k(a.k().max(b.k()));
 52    }
 53
 54    fn add_inplace<A>(&mut self, module: &Module<FFT64>, a: &A)
 55    where
 56        A: GLWECiphertextToRef + Infos,
 57    {
 58        #[cfg(debug_assertions)]
 59        {
 60            assert_eq!(a.n(), module.n());
 61            assert_eq!(self.n(), module.n());
 62            assert_eq!(self.basek(), a.basek());
 63            assert!(self.rank() >= a.rank())
 64        }
 65
 66        let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
 67        let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
 68
 69        (0..a.rank() + 1).for_each(|i| {
 70            module.vec_znx_add_inplace(&mut self_mut.data, i, &a_ref.data, i);
 71        });
 72
 73        self.set_k(a.k().max(self.k()));
 74    }
 75
 76    fn sub<A, B>(&mut self, module: &Module<FFT64>, a: &A, b: &B)
 77    where
 78        A: GLWECiphertextToRef,
 79        B: GLWECiphertextToRef,
 80    {
 81        #[cfg(debug_assertions)]
 82        {
 83            assert_eq!(a.n(), module.n());
 84            assert_eq!(b.n(), module.n());
 85            assert_eq!(self.n(), module.n());
 86            assert_eq!(a.basek(), b.basek());
 87            assert!(self.rank() >= a.rank().max(b.rank()));
 88        }
 89
 90        let min_col: usize = a.rank().min(b.rank()) + 1;
 91        let max_col: usize = a.rank().max(b.rank() + 1);
 92        let self_col: usize = self.rank() + 1;
 93
 94        let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
 95        let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
 96        let b_ref: &GLWECiphertext<&[u8]> = &b.to_ref();
 97
 98        (0..min_col).for_each(|i| {
 99            module.vec_znx_sub(&mut self_mut.data, i, &a_ref.data, i, &b_ref.data, i);
100        });
101
102        if a.rank() > b.rank() {
103            (min_col..max_col).for_each(|i| {
104                module.vec_znx_copy(&mut self_mut.data, i, &a_ref.data, i);
105            });
106        } else {
107            (min_col..max_col).for_each(|i| {
108                module.vec_znx_copy(&mut self_mut.data, i, &b_ref.data, i);
109                module.vec_znx_negate_inplace(&mut self_mut.data, i);
110            });
111        }
112
113        let size: usize = self_mut.size();
114        (max_col..self_col).for_each(|i| {
115            (0..size).for_each(|j| {
116                self_mut.data.zero_at(i, j);
117            });
118        });
119
120        self.set_basek(a.basek());
121        self.set_k(a.k().max(b.k()));
122    }
123
124    fn sub_inplace_ab<A>(&mut self, module: &Module<FFT64>, a: &A)
125    where
126        A: GLWECiphertextToRef + Infos,
127    {
128        #[cfg(debug_assertions)]
129        {
130            assert_eq!(a.n(), module.n());
131            assert_eq!(self.n(), module.n());
132            assert_eq!(self.basek(), a.basek());
133            assert!(self.rank() >= a.rank())
134        }
135
136        let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
137        let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
138
139        (0..a.rank() + 1).for_each(|i| {
140            module.vec_znx_sub_ab_inplace(&mut self_mut.data, i, &a_ref.data, i);
141        });
142
143        self.set_k(a.k().max(self.k()));
144    }
145
146    fn sub_inplace_ba<A>(&mut self, module: &Module<FFT64>, a: &A)
147    where
148        A: GLWECiphertextToRef + Infos,
149    {
150        #[cfg(debug_assertions)]
151        {
152            assert_eq!(a.n(), module.n());
153            assert_eq!(self.n(), module.n());
154            assert_eq!(self.basek(), a.basek());
155            assert!(self.rank() >= a.rank())
156        }
157
158        let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
159        let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
160
161        (0..a.rank() + 1).for_each(|i| {
162            module.vec_znx_sub_ba_inplace(&mut self_mut.data, i, &a_ref.data, i);
163        });
164
165        self.set_k(a.k().max(self.k()));
166    }
167
168    fn rotate<A>(&mut self, module: &Module<FFT64>, k: i64, a: &A)
169    where
170        A: GLWECiphertextToRef + Infos,
171    {
172        #[cfg(debug_assertions)]
173        {
174            assert_eq!(a.n(), module.n());
175            assert_eq!(self.n(), module.n());
176            assert_eq!(self.rank(), a.rank())
177        }
178
179        let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
180        let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
181
182        (0..a.rank() + 1).for_each(|i| {
183            module.vec_znx_rotate(k, &mut self_mut.data, i, &a_ref.data, i);
184        });
185
186        self.set_basek(a.basek());
187        self.set_k(a.k());
188    }
189
190    fn rotate_inplace(&mut self, module: &Module<FFT64>, k: i64) {
191        #[cfg(debug_assertions)]
192        {
193            assert_eq!(self.n(), module.n());
194        }
195
196        let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
197
198        (0..self_mut.rank() + 1).for_each(|i| {
199            module.vec_znx_rotate_inplace(k, &mut self_mut.data, i);
200        });
201    }
202
203    fn copy<A>(&mut self, module: &Module<FFT64>, a: &A)
204    where
205        A: GLWECiphertextToRef + Infos,
206    {
207        #[cfg(debug_assertions)]
208        {
209            assert_eq!(self.n(), module.n());
210            assert_eq!(a.n(), module.n());
211            assert_eq!(self.rank(), a.rank());
212        }
213
214        let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
215        let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
216
217        (0..self_mut.rank() + 1).for_each(|i| {
218            module.vec_znx_copy(&mut self_mut.data, i, &a_ref.data, i);
219        });
220
221        self.set_k(a.k());
222        self.set_basek(a.basek());
223    }
224
225    fn rsh(&mut self, k: usize, scratch: &mut Scratch) {
226        let basek: usize = self.basek();
227        let mut self_mut: GLWECiphertext<&mut [u8]> = self.to_mut();
228        self_mut.data.rsh(basek, k, scratch);
229    }
230
231    fn normalize<A>(&mut self, module: &Module<FFT64>, a: &A, scratch: &mut Scratch)
232    where
233        A: GLWECiphertextToRef,
234    {
235        #[cfg(debug_assertions)]
236        {
237            assert_eq!(self.n(), module.n());
238            assert_eq!(a.n(), module.n());
239            assert_eq!(self.rank(), a.rank());
240        }
241
242        let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
243        let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
244
245        (0..self_mut.rank() + 1).for_each(|i| {
246            module.vec_znx_normalize(a.basek(), &mut self_mut.data, i, &a_ref.data, i, scratch);
247        });
248        self.set_basek(a.basek());
249        self.set_k(a.k());
250    }
251
252    fn normalize_inplace(&mut self, module: &Module<FFT64>, scratch: &mut Scratch) {
253        #[cfg(debug_assertions)]
254        {
255            assert_eq!(self.n(), module.n());
256        }
257        let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
258        (0..self_mut.rank() + 1).for_each(|i| {
259            module.vec_znx_normalize_inplace(self_mut.basek(), &mut self_mut.data, i, scratch);
260        });
261    }
262}
263
264impl GLWECiphertext<Vec<u8>> {
265    pub fn rsh_scratch_space(module: &Module<FFT64>) -> usize {
266        VecZnx::rsh_scratch_space(module.n())
267    }
268}