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    fn elwise_fma_scalar_mut(
 16        &self,
 17        a: &mut [Self::Element],
 18        b: &[Self::Element],
 19        c: &Self::Element,
 20    );
 21
 22    fn modulus(&self) -> Self::Element;
 23}
 24
 25pub trait ArithmeticOps {
 26    type Element;
 27
 28    fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
 29    fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
 30    fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
 31    fn neg(&self, a: &Self::Element) -> Self::Element;
 32
 33    fn modulus(&self) -> Self::Element;
 34}
 35
 36pub struct ModularOpsU64 {
 37    q: u64,
 38    logq: usize,
 39    barrett_mu: u128,
 40    barrett_alpha: usize,
 41}
 42
 43impl ModularOpsU64 {
 44    pub fn new(q: u64) -> ModularOpsU64 {
 45        let logq = 64 - q.leading_zeros();
 46
 47        // barrett calculation
 48        let mu = (1u128 << (logq * 2 + 3)) / (q as u128);
 49        let alpha = logq + 3;
 50
 51        ModularOpsU64 {
 52            q,
 53            logq: logq as usize,
 54            barrett_alpha: alpha as usize,
 55            barrett_mu: mu,
 56        }
 57    }
 58
 59    fn add_mod_fast(&self, a: u64, b: u64) -> u64 {
 60        debug_assert!(a < self.q);
 61        debug_assert!(b < self.q);
 62
 63        let mut o = a + b;
 64        if o >= self.q {
 65            o -= self.q;
 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 , q)
 85    fn mul_mod_fast(&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        let mut out = (ab - tmp) as u64;
102
103        if out >= self.q {
104            out -= self.q;
105        }
106
107        return out;
108    }
109}
110
111impl ArithmeticOps for ModularOpsU64 {
112    type Element = u64;
113
114    fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
115        self.add_mod_fast(*a, *b)
116    }
117
118    fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
119        self.mul_mod_fast(*a, *b)
120    }
121
122    fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
123        self.sub_mod_fast(*a, *b)
124    }
125
126    fn neg(&self, a: &Self::Element) -> Self::Element {
127        self.q - *a
128    }
129
130    fn modulus(&self) -> Self::Element {
131        self.q
132    }
133}
134
135impl VectorOps for ModularOpsU64 {
136    type Element = u64;
137
138    fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
139        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
140            *ai = self.add_mod_fast(*ai, *bi);
141        });
142    }
143
144    fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
145        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
146            *ai = self.sub_mod_fast(*ai, *bi);
147        });
148    }
149
150    fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
151        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
152            *ai = self.mul_mod_fast(*ai, *bi);
153        });
154    }
155
156    fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
157        a.iter_mut().for_each(|ai| *ai = self.q - *ai);
158    }
159
160    fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
161        izip!(out.iter_mut(), a.iter()).for_each(|(oi, ai)| {
162            *oi = self.mul_mod_fast(*ai, *b);
163        });
164    }
165
166    fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]) {
167        izip!(out.iter_mut(), a.iter(), b.iter()).for_each(|(oi, ai, bi)| {
168            *oi = self.mul_mod_fast(*ai, *bi);
169        });
170    }
171
172    fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
173        izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(ai, bi, ci)| {
174            *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *ci));
175        });
176    }
177
178    fn elwise_fma_scalar_mut(
179        &self,
180        a: &mut [Self::Element],
181        b: &[Self::Element],
182        c: &Self::Element,
183    ) {
184        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
185            *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *c));
186        });
187    }
188
189    fn modulus(&self) -> Self::Element {
190        self.q
191    }
192}