phantom-zone/src/backend.rs
src/backend.rsRaw
1use std::marker::PhantomData;
2
3use itertools::izip;
4use num_traits::{PrimInt, Signed, ToPrimitive, WrappingAdd, WrappingMul, WrappingSub, Zero};
5
6pub trait Modulus {
7 type Element;
8 /// Modulus value if it fits in Element
9 fn q(&self) -> Option<Self::Element>;
10 /// Is modulus native?
11 fn is_native(&self) -> bool;
12 /// -1 in signed representaiton
13 fn neg_one(&self) -> Self::Element;
14 /// Largest unsigned value that fits in the modulus. That is, q - 1.
15 fn largest_unsigned_value(&self) -> Self::Element;
16 /// Smallest unsigned value that fits in the modulus
17 /// Always assmed to be 0.
18 fn smallest_unsigned_value(&self) -> Self::Element;
19 /// Convert unsigned value in signed represetation to i64
20 fn to_i64(&self, v: &Self::Element) -> i64;
21 /// Convert f64 to signed represented in modulus
22 fn from_f64(&self, v: f64) -> Self::Element;
23 /// Convert i64 to signed represented in modulus
24 fn from_i64(&self, v: i64) -> Self::Element;
25}
26
27impl Modulus for u64 {
28 type Element = u64;
29 fn is_native(&self) -> bool {
30 // q of size u64 can never be a naitve modulus
31 false
32 }
33 fn largest_unsigned_value(&self) -> Self::Element {
34 self - 1
35 }
36 fn neg_one(&self) -> Self::Element {
37 self - 1
38 }
39 fn smallest_unsigned_value(&self) -> Self::Element {
40 0
41 }
42 fn to_i64(&self, v: &Self::Element) -> i64 {
43 assert!(v < self);
44
45 if *v > (self >> 1) {
46 ToPrimitive::to_i64(&(self - v)).unwrap()
47 } else {
48 ToPrimitive::to_i64(v).unwrap()
49 }
50 }
51 fn from_f64(&self, v: f64) -> Self::Element {
52 //FIXME (Jay): Before I check whether v is smaller than 0 with `let is_neg =
53 // o.is_sign_negative() && o != 0.0; I'm ocnfused why didn't I simply check <
54 // 0.0?
55 let v = v.round();
56 if v < 0.0 {
57 self - v.to_u64().unwrap()
58 } else {
59 v.to_u64().unwrap()
60 }
61 }
62 fn from_i64(&self, v: i64) -> Self::Element {
63 if v < 0 {
64 self - v.to_u64().unwrap()
65 } else {
66 v.to_u64().unwrap()
67 }
68 }
69 fn q(&self) -> Option<Self::Element> {
70 Some(*self)
71 }
72}
73
74pub trait ModInit {
75 type M;
76 fn new(modulus: Self::M) -> Self;
77}
78
79pub trait GetModulus {
80 type Element;
81 type M: Modulus<Element = Self::Element>;
82 fn modulus(&self) -> &Self::M;
83}
84
85pub trait VectorOps {
86 type Element;
87
88 fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element);
89 fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]);
90
91 fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
92 fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
93 fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
94 fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element);
95 fn elwise_neg_mut(&self, a: &mut [Self::Element]);
96 /// inplace mutates `a`: a = a + b*c
97 fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]);
98 fn elwise_fma_scalar_mut(
99 &self,
100 a: &mut [Self::Element],
101 b: &[Self::Element],
102 c: &Self::Element,
103 );
104
105 // fn modulus(&self) -> Self::Element;
106}
107
108pub trait ArithmeticOps {
109 type Element;
110
111 fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
112 fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
113 fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
114 fn neg(&self, a: &Self::Element) -> Self::Element;
115
116 // fn modulus(&self) -> Self::Element;
117}
118
119pub struct ModularOpsU64<T> {
120 q: u64,
121 logq: usize,
122 barrett_mu: u128,
123 barrett_alpha: usize,
124 modulus: T,
125}
126
127impl<T> ModInit for ModularOpsU64<T>
128where
129 T: Modulus<Element = u64>,
130{
131 type M = T;
132 fn new(modulus: Self::M) -> ModularOpsU64<T> {
133 assert!(!modulus.is_native());
134
135 // largest unsigned value modulus fits is modulus-1
136 let q = modulus.largest_unsigned_value() + 1;
137 let logq = 64 - (q + 1u64).leading_zeros();
138
139 // barrett calculation
140 let mu = (1u128 << (logq * 2 + 3)) / (q as u128);
141 let alpha = logq + 3;
142
143 ModularOpsU64 {
144 q,
145 logq: logq as usize,
146 barrett_alpha: alpha as usize,
147 barrett_mu: mu,
148 modulus,
149 }
150 }
151}
152
153impl<T> ModularOpsU64<T> {
154 fn add_mod_fast(&self, a: u64, b: u64) -> u64 {
155 debug_assert!(a < self.q);
156 debug_assert!(b < self.q);
157
158 let mut o = a + b;
159 if o >= self.q {
160 o -= self.q;
161 }
162 o
163 }
164
165 fn sub_mod_fast(&self, a: u64, b: u64) -> u64 {
166 debug_assert!(a < self.q);
167 debug_assert!(b < self.q);
168
169 if a > b {
170 a - b
171 } else {
172 (self.q + a) - b
173 }
174 }
175
176 /// returns (a * b) % q
177 ///
178 /// - both a and b must be in range [0, 2q)
179 /// - output is in range [0 , q)
180 fn mul_mod_fast(&self, a: u64, b: u64) -> u64 {
181 debug_assert!(a < 2 * self.q);
182 debug_assert!(b < 2 * self.q);
183
184 let ab = a as u128 * b as u128;
185
186 // ab / (2^{n + \beta})
187 // note: \beta is assumed to -2
188 let tmp = ab >> (self.logq - 2);
189
190 // k = ((ab / (2^{n + \beta})) * \mu) / 2^{\alpha - (-2)}
191 let k = (tmp * self.barrett_mu) >> (self.barrett_alpha + 2);
192
193 // ab - k*p
194 let tmp = k * (self.q as u128);
195
196 let mut out = (ab - tmp) as u64;
197
198 if out >= self.q {
199 out -= self.q;
200 }
201
202 return out;
203 }
204}
205
206impl<T> ArithmeticOps for ModularOpsU64<T> {
207 type Element = u64;
208
209 fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
210 self.add_mod_fast(*a, *b)
211 }
212
213 fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
214 self.mul_mod_fast(*a, *b)
215 }
216
217 fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
218 self.sub_mod_fast(*a, *b)
219 }
220
221 fn neg(&self, a: &Self::Element) -> Self::Element {
222 self.q - *a
223 }
224
225 // fn modulus(&self) -> Self::Element {
226 // self.q
227 // }
228}
229
230impl<T> VectorOps for ModularOpsU64<T> {
231 type Element = u64;
232
233 fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
234 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
235 *ai = self.add_mod_fast(*ai, *bi);
236 });
237 }
238
239 fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
240 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
241 *ai = self.sub_mod_fast(*ai, *bi);
242 });
243 }
244
245 fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
246 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
247 *ai = self.mul_mod_fast(*ai, *bi);
248 });
249 }
250
251 fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
252 a.iter_mut().for_each(|ai| *ai = self.q - *ai);
253 }
254
255 fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
256 izip!(out.iter_mut(), a.iter()).for_each(|(oi, ai)| {
257 *oi = self.mul_mod_fast(*ai, *b);
258 });
259 }
260
261 fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]) {
262 izip!(out.iter_mut(), a.iter(), b.iter()).for_each(|(oi, ai, bi)| {
263 *oi = self.mul_mod_fast(*ai, *bi);
264 });
265 }
266
267 fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element) {
268 a.iter_mut().for_each(|ai| {
269 *ai = self.mul_mod_fast(*ai, *b);
270 });
271 }
272
273 fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
274 izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(ai, bi, ci)| {
275 *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *ci));
276 });
277 }
278
279 fn elwise_fma_scalar_mut(
280 &self,
281 a: &mut [Self::Element],
282 b: &[Self::Element],
283 c: &Self::Element,
284 ) {
285 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
286 *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *c));
287 });
288 }
289
290 // fn modulus(&self) -> Self::Element {
291 // self.q
292 // }
293}
294
295impl<T> GetModulus for ModularOpsU64<T>
296where
297 T: Modulus,
298{
299 type Element = T::Element;
300 type M = T;
301 fn modulus(&self) -> &Self::M {
302 &self.modulus
303 }
304}
305
306pub struct WordSizeModulus<T> {
307 modulus: T,
308}
309
310impl<T> ModInit for WordSizeModulus<T>
311where
312 T: Modulus,
313{
314 type M = T;
315 fn new(modulus: T) -> Self {
316 assert!(modulus.is_native());
317 // For now assume ModulusOpsU64 is only used for u64
318 Self { modulus: modulus }
319 }
320}
321
322impl<T> ArithmeticOps for WordSizeModulus<T>
323where
324 T: Modulus,
325 T::Element: WrappingAdd + WrappingSub + WrappingMul + Zero,
326{
327 type Element = T::Element;
328 fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
329 T::Element::wrapping_add(a, b)
330 }
331
332 fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
333 T::Element::wrapping_mul(a, b)
334 }
335
336 fn neg(&self, a: &Self::Element) -> Self::Element {
337 T::Element::wrapping_sub(&T::Element::zero(), a)
338 }
339
340 fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
341 T::Element::wrapping_sub(a, b)
342 }
343
344 // fn modulus(&self) -> &T {
345 // &self.modulus
346 // }
347}
348
349impl<T> VectorOps for WordSizeModulus<T>
350where
351 T: Modulus,
352 T::Element: WrappingAdd + WrappingSub + WrappingMul + Zero,
353{
354 type Element = T::Element;
355
356 fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
357 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
358 *ai = T::Element::wrapping_add(ai, bi);
359 });
360 }
361
362 fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
363 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
364 *ai = T::Element::wrapping_sub(ai, bi);
365 });
366 }
367
368 fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
369 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
370 *ai = T::Element::wrapping_mul(ai, bi);
371 });
372 }
373
374 fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
375 a.iter_mut()
376 .for_each(|ai| *ai = T::Element::wrapping_sub(&T::Element::zero(), ai));
377 }
378
379 fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
380 izip!(out.iter_mut(), a.iter()).for_each(|(oi, ai)| {
381 *oi = T::Element::wrapping_mul(ai, b);
382 });
383 }
384
385 fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]) {
386 izip!(out.iter_mut(), a.iter(), b.iter()).for_each(|(oi, ai, bi)| {
387 *oi = T::Element::wrapping_mul(ai, bi);
388 });
389 }
390
391 fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element) {
392 a.iter_mut().for_each(|ai| {
393 *ai = T::Element::wrapping_mul(ai, b);
394 });
395 }
396
397 fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
398 izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(ai, bi, ci)| {
399 *ai = T::Element::wrapping_add(ai, &T::Element::wrapping_mul(bi, ci));
400 });
401 }
402
403 fn elwise_fma_scalar_mut(
404 &self,
405 a: &mut [Self::Element],
406 b: &[Self::Element],
407 c: &Self::Element,
408 ) {
409 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
410 *ai = T::Element::wrapping_add(ai, &T::Element::wrapping_mul(bi, c));
411 });
412 }
413
414 // fn modulus(&self) -> &T {
415 // &self.modulus
416 // }
417}