1use itertools::izip;
  2use num_traits::WrappingMul;
  3
  4use super::{
  5    ArithmeticLazyOps, ArithmeticOps, GetModulus, ModInit, Modulus, ShoupMatrixFMA, VectorOps,
  6};
  7use crate::RowMut;
  8
  9pub struct ModularOpsU64<T> {
 10    q: u64,
 11    q_twice: u64,
 12    logq: usize,
 13    barrett_mu: u128,
 14    barrett_alpha: usize,
 15    modulus: T,
 16}
 17
 18impl<T> ModInit for ModularOpsU64<T>
 19where
 20    T: Modulus<Element = u64>,
 21{
 22    type M = T;
 23    fn new(modulus: Self::M) -> ModularOpsU64<T> {
 24        assert!(!modulus.is_native());
 25
 26        // largest unsigned value modulus fits is modulus-1
 27        let q = modulus.largest_unsigned_value() + 1;
 28        let logq = 64 - (q + 1u64).leading_zeros();
 29
 30        // barrett calculation
 31        let mu = (1u128 << (logq * 2 + 3)) / (q as u128);
 32        let alpha = logq + 3;
 33
 34        ModularOpsU64 {
 35            q,
 36            q_twice: q << 1,
 37            logq: logq as usize,
 38            barrett_alpha: alpha as usize,
 39            barrett_mu: mu,
 40            modulus,
 41        }
 42    }
 43}
 44
 45impl<T> ModularOpsU64<T> {
 46    fn add_mod_fast(&self, a: u64, b: u64) -> u64 {
 47        debug_assert!(a < self.q);
 48        debug_assert!(b < self.q);
 49
 50        let mut o = a + b;
 51        if o >= self.q {
 52            o -= self.q;
 53        }
 54        o
 55    }
 56
 57    fn add_mod_fast_lazy(&self, a: u64, b: u64) -> u64 {
 58        debug_assert!(a < self.q_twice);
 59        debug_assert!(b < self.q_twice);
 60
 61        let mut o = a + b;
 62        if o >= self.q_twice {
 63            o -= self.q_twice;
 64        }
 65        o
 66    }
 67
 68    fn sub_mod_fast(&self, a: u64, b: u64) -> u64 {
 69        debug_assert!(a < self.q);
 70        debug_assert!(b < self.q);
 71
 72        if a >= b {
 73            a - b
 74        } else {
 75            (self.q + a) - b
 76        }
 77    }
 78
 79    // returns (a * b)  % q
 80    ///
 81    /// - both a and b must be in range [0, 2q)
 82    /// - output is in range [0 , 2q)
 83    fn mul_mod_fast_lazy(&self, a: u64, b: u64) -> u64 {
 84        debug_assert!(a < 2 * self.q);
 85        debug_assert!(b < 2 * self.q);
 86
 87        let ab = a as u128 * b as u128;
 88
 89        // ab / (2^{n + \beta})
 90        // note: \beta is assumed to -2
 91        let tmp = ab >> (self.logq - 2);
 92
 93        // k = ((ab / (2^{n + \beta})) * \mu) / 2^{\alpha - (-2)}
 94        let k = (tmp * self.barrett_mu) >> (self.barrett_alpha + 2);
 95
 96        // ab - k*p
 97        let tmp = k * (self.q as u128);
 98
 99        (ab - tmp) as u64
100    }
101
102    /// returns (a * b)  % q
103    ///
104    /// - both a and b must be in range [0, 2q)
105    /// - output is in range [0 , q)
106    fn mul_mod_fast(&self, a: u64, b: u64) -> u64 {
107        debug_assert!(a < 2 * self.q);
108        debug_assert!(b < 2 * self.q);
109
110        let ab = a as u128 * b as u128;
111
112        // ab / (2^{n + \beta})
113        // note: \beta is assumed to -2
114        let tmp = ab >> (self.logq - 2);
115
116        // k = ((ab / (2^{n + \beta})) * \mu) / 2^{\alpha - (-2)}
117        let k = (tmp * self.barrett_mu) >> (self.barrett_alpha + 2);
118
119        // ab - k*p
120        let tmp = k * (self.q as u128);
121
122        let mut out = (ab - tmp) as u64;
123
124        if out >= self.q {
125            out -= self.q;
126        }
127
128        return out;
129    }
130}
131
132impl<T> ArithmeticOps for ModularOpsU64<T> {
133    type Element = u64;
134
135    fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
136        self.add_mod_fast(*a, *b)
137    }
138
139    fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
140        self.mul_mod_fast(*a, *b)
141    }
142
143    fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
144        self.sub_mod_fast(*a, *b)
145    }
146
147    fn neg(&self, a: &Self::Element) -> Self::Element {
148        self.q - *a
149    }
150
151    // fn modulus(&self) -> Self::Element {
152    //     self.q
153    // }
154}
155
156impl<T> ArithmeticLazyOps for ModularOpsU64<T> {
157    type Element = u64;
158    fn add_lazy(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
159        self.add_mod_fast_lazy(*a, *b)
160    }
161    fn mul_lazy(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
162        self.mul_mod_fast_lazy(*a, *b)
163    }
164}
165
166impl<T> VectorOps for ModularOpsU64<T> {
167    type Element = u64;
168
169    fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
170        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
171            *ai = self.add_mod_fast(*ai, *bi);
172        });
173    }
174
175    fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
176        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
177            *ai = self.sub_mod_fast(*ai, *bi);
178        });
179    }
180
181    fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
182        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
183            *ai = self.mul_mod_fast(*ai, *bi);
184        });
185    }
186
187    fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
188        a.iter_mut().for_each(|ai| *ai = self.q - *ai);
189    }
190
191    fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
192        izip!(out.iter_mut(), a.iter()).for_each(|(oi, ai)| {
193            *oi = self.mul_mod_fast(*ai, *b);
194        });
195    }
196
197    fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]) {
198        izip!(out.iter_mut(), a.iter(), b.iter()).for_each(|(oi, ai, bi)| {
199            *oi = self.mul_mod_fast(*ai, *bi);
200        });
201    }
202
203    fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element) {
204        a.iter_mut().for_each(|ai| {
205            *ai = self.mul_mod_fast(*ai, *b);
206        });
207    }
208
209    fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
210        izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(ai, bi, ci)| {
211            *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *ci));
212        });
213    }
214
215    fn elwise_fma_scalar_mut(
216        &self,
217        a: &mut [Self::Element],
218        b: &[Self::Element],
219        c: &Self::Element,
220    ) {
221        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
222            *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *c));
223        });
224    }
225
226    // fn modulus(&self) -> Self::Element {
227    //     self.q
228    // }
229}
230
231impl<R: RowMut<Element = u64>, T> ShoupMatrixFMA<R> for ModularOpsU64<T> {
232    fn shoup_matrix_fma(&self, out: &mut [R::Element], a: &[R], a_shoup: &[R], b: &[R]) {
233        assert!(a.len() == a_shoup.len());
234        assert!(a.len() == b.len());
235
236        let q = self.q;
237        let q_twice = self.q << 1;
238
239        izip!(a.iter(), a_shoup.iter(), b.iter()).for_each(|(a_row, a_shoup_row, b_row)| {
240            izip!(
241                out.as_mut().iter_mut(),
242                a_row.as_ref().iter(),
243                a_shoup_row.as_ref().iter(),
244                b_row.as_ref().iter()
245            )
246            .for_each(|(o, a0, a0_shoup, b0)| {
247                let quotient = ((*a0_shoup as u128 * *b0 as u128) >> 64) as u64;
248                let mut v = (a0.wrapping_mul(b0)).wrapping_add(*o);
249                v = v.wrapping_sub(q.wrapping_mul(quotient));
250
251                if v >= q_twice {
252                    v -= q_twice;
253                }
254
255                *o = v;
256            });
257        });
258    }
259}
260
261impl<T> GetModulus for ModularOpsU64<T>
262where
263    T: Modulus,
264{
265    type Element = T::Element;
266    type M = T;
267    fn modulus(&self) -> &Self::M {
268        &self.modulus
269    }
270}
271
272#[cfg(test)]
273mod tests {
274    use super::*;
275    use itertools::Itertools;
276    use rand::{thread_rng, Rng};
277    use rand_distr::Uniform;
278
279    #[test]
280    fn fma() {
281        let mut rng = thread_rng();
282        let prime = 36028797017456641;
283        let ring_size = 1 << 3;
284
285        let dist = Uniform::new(0, prime);
286        let d = 2;
287        let a0_matrix = (0..d)
288            .into_iter()
289            .map(|_| (&mut rng).sample_iter(dist).take(ring_size).collect_vec())
290            .collect_vec();
291        // a0 in shoup representation
292        let a0_shoup_matrix = a0_matrix
293            .iter()
294            .map(|r| {
295                r.iter()
296                    .map(|v| {
297                        // $(v * 2^{\beta}) / p$
298                        ((*v as u128 * (1u128 << 64)) / prime as u128) as u64
299                    })
300                    .collect_vec()
301            })
302            .collect_vec();
303        let a1_matrix = (0..d)
304            .into_iter()
305            .map(|_| (&mut rng).sample_iter(dist).take(ring_size).collect_vec())
306            .collect_vec();
307
308        let modop = ModularOpsU64::new(prime);
309
310        let mut out_shoup_fma_lazy = vec![0u64; ring_size];
311        modop.shoup_matrix_fma(
312            &mut out_shoup_fma_lazy,
313            &a0_matrix,
314            &a0_shoup_matrix,
315            &a1_matrix,
316        );
317        let out_shoup_fma = out_shoup_fma_lazy
318            .iter()
319            .map(|v| if *v >= prime { v - prime } else { *v })
320            .collect_vec();
321
322        // expected
323        let mut out_expected = vec![0u64; ring_size];
324        izip!(a0_matrix.iter(), a1_matrix.iter()).for_each(|(a_r, b_r)| {
325            izip!(out_expected.iter_mut(), a_r.iter(), b_r.iter()).for_each(|(o, a0, a1)| {
326                *o = (*o + ((*a0 as u128 * *a1 as u128) % prime as u128) as u64) % prime;
327            });
328        });
329
330        assert_eq!(out_expected, out_shoup_fma);
331    }
332}