1use itertools::izip;
  2
  3pub trait VectorOps {
  4    type Element;
  5
  6    fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element);
  7    fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]);
  8
  9    fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
 10    fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
 11    fn elwise_neg_mut(&self, a: &mut [Self::Element]);
 12    /// inplace mutates `a`: a = a + b*c
 13    fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]);
 14
 15    fn modulus(&self) -> Self::Element;
 16}
 17
 18pub trait ArithmeticOps {
 19    type Element;
 20
 21    fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
 22    fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
 23    fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
 24
 25    fn modulus(&self) -> Self::Element;
 26}
 27
 28pub struct ModularOpsU64 {
 29    q: u64,
 30    logq: usize,
 31    barrett_mu: u128,
 32    barrett_alpha: usize,
 33}
 34
 35impl ModularOpsU64 {
 36    pub fn new(q: u64) -> ModularOpsU64 {
 37        let logq = 64 - q.leading_zeros();
 38
 39        // barrett calculation
 40        let mu = (1u128 << (logq * 2 + 3)) / (q as u128);
 41        let alpha = logq + 3;
 42
 43        ModularOpsU64 {
 44            q,
 45            logq: logq as usize,
 46            barrett_alpha: alpha as usize,
 47            barrett_mu: mu,
 48        }
 49    }
 50
 51    fn add_mod_fast(&self, a: u64, b: u64) -> u64 {
 52        debug_assert!(a < self.q);
 53        debug_assert!(b < self.q);
 54
 55        let mut o = a + b;
 56        if o >= self.q {
 57            o -= self.q;
 58        }
 59        o
 60    }
 61
 62    fn sub_mod_fast(&self, a: u64, b: u64) -> u64 {
 63        debug_assert!(a < self.q);
 64        debug_assert!(b < self.q);
 65
 66        if a > b {
 67            a - b
 68        } else {
 69            (self.q + a) - b
 70        }
 71    }
 72
 73    /// returns (a * b)  % q
 74    ///
 75    /// - both a and b must be in range [0, 2q)
 76    /// - output is in range [0 , q)
 77    fn mul_mod_fast(&self, a: u64, b: u64) -> u64 {
 78        debug_assert!(a < 2 * self.q);
 79        debug_assert!(b < 2 * self.q);
 80
 81        let ab = a as u128 * b as u128;
 82
 83        // ab / (2^{n + \beta})
 84        // note: \beta is assumed to -2
 85        let tmp = ab >> (self.logq - 2);
 86
 87        // k = ((ab / (2^{n + \beta})) * \mu) / 2^{\alpha - (-2)}
 88        let k = (tmp * self.barrett_mu) >> (self.barrett_alpha + 2);
 89
 90        // ab - k*p
 91        let tmp = k * (self.q as u128);
 92
 93        let mut out = (ab - tmp) as u64;
 94
 95        if out >= self.q {
 96            out -= self.q;
 97        }
 98
 99        return out;
100    }
101}
102
103impl ArithmeticOps for ModularOpsU64 {
104    type Element = u64;
105
106    fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
107        self.add_mod_fast(*a, *b)
108    }
109
110    fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
111        self.mul_mod_fast(*a, *b)
112    }
113
114    fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
115        self.sub_mod_fast(*a, *b)
116    }
117
118    fn modulus(&self) -> Self::Element {
119        self.q
120    }
121}
122
123impl VectorOps for ModularOpsU64 {
124    type Element = u64;
125
126    fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
127        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
128            *ai = self.add_mod_fast(*ai, *bi);
129        });
130    }
131
132    fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
133        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
134            *ai = self.mul_mod_fast(*ai, *bi);
135        });
136    }
137
138    fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
139        a.iter_mut().for_each(|ai| *ai = self.q - *ai);
140    }
141
142    fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
143        izip!(out.iter_mut(), a.iter()).for_each(|(oi, ai)| {
144            *oi = self.mul_mod_fast(*ai, *b);
145        });
146    }
147
148    fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]) {
149        izip!(out.iter_mut(), a.iter(), b.iter()).for_each(|(oi, ai, bi)| {
150            *oi = self.mul_mod_fast(*ai, *bi);
151        });
152    }
153
154    fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
155        izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(ai, bi, ci)| {
156            *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *ci));
157        });
158    }
159
160    fn modulus(&self) -> Self::Element {
161        self.q
162    }
163}