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