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