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