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<R: RowMut<Element = u64>, T> ShoupMatrixFMA<R> for ModularOpsU64<T> {
234    fn shoup_matrix_fma(&self, out: &mut [R::Element], a: &[R], a_shoup: &[R], b: &[R]) {
235        assert!(a.len() == a_shoup.len());
236        assert!(a.len() == b.len());
237
238        let q = self.q;
239        let q_twice = self.q << 1;
240
241        izip!(a.iter(), a_shoup.iter(), b.iter()).for_each(|(a_row, a_shoup_row, b_row)| {
242            izip!(
243                out.as_mut().iter_mut(),
244                a_row.as_ref().iter(),
245                a_shoup_row.as_ref().iter(),
246                b_row.as_ref().iter()
247            )
248            .for_each(|(o, a0, a0_shoup, b0)| {
249                let quotient = ((*a0_shoup as u128 * *b0 as u128) >> 64) as u64;
250                let mut v = (a0.wrapping_mul(b0)).wrapping_add(*o);
251                v = v.wrapping_sub(q.wrapping_mul(quotient));
252
253                if v >= q_twice {
254                    v -= q_twice;
255                }
256
257                *o = v;
258            });
259        });
260    }
261}
262
263impl<T> GetModulus for ModularOpsU64<T>
264where
265    T: Modulus,
266{
267    type Element = T::Element;
268    type M = T;
269    fn modulus(&self) -> &Self::M {
270        &self.modulus
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use super::*;
277    use itertools::Itertools;
278    use rand::{thread_rng, Rng};
279    use rand_distr::Uniform;
280
281    #[test]
282    fn fma() {
283        let mut rng = thread_rng();
284        let prime = 36028797017456641;
285        let ring_size = 1 << 3;
286
287        let dist = Uniform::new(0, prime);
288        let d = 2;
289        let a0_matrix = (0..d)
290            .into_iter()
291            .map(|_| (&mut rng).sample_iter(dist).take(ring_size).collect_vec())
292            .collect_vec();
293        // a0 in shoup representation
294        let a0_shoup_matrix = a0_matrix
295            .iter()
296            .map(|r| {
297                r.iter()
298                    .map(|v| {
299                        // $(v * 2^{\beta}) / p$
300                        ((*v as u128 * (1u128 << 64)) / prime as u128) as u64
301                    })
302                    .collect_vec()
303            })
304            .collect_vec();
305        let a1_matrix = (0..d)
306            .into_iter()
307            .map(|_| (&mut rng).sample_iter(dist).take(ring_size).collect_vec())
308            .collect_vec();
309
310        let modop = ModularOpsU64::new(prime);
311
312        let mut out_shoup_fma_lazy = vec![0u64; ring_size];
313        modop.shoup_matrix_fma(
314            &mut out_shoup_fma_lazy,
315            &a0_matrix,
316            &a0_shoup_matrix,
317            &a1_matrix,
318        );
319        let out_shoup_fma = out_shoup_fma_lazy
320            .iter()
321            .map(|v| if *v >= prime { v - prime } else { *v })
322            .collect_vec();
323
324        // expected
325        let mut out_expected = vec![0u64; ring_size];
326        izip!(a0_matrix.iter(), a1_matrix.iter()).for_each(|(a_r, b_r)| {
327            izip!(out_expected.iter_mut(), a_r.iter(), b_r.iter()).for_each(|(o, a0, a1)| {
328                *o = (*o + ((*a0 as u128 * *a1 as u128) % prime as u128) as u64) % prime;
329            });
330        });
331
332        assert_eq!(out_expected, out_shoup_fma);
333    }
334}