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