1use std::marker::PhantomData;
  2
  3use itertools::izip;
  4use num_traits::{WrappingAdd, WrappingMul, WrappingSub};
  5
  6pub trait Modulus {
  7    type Element;
  8    fn is_native() -> bool;
  9    /// -1 in signed representaiton
 10    fn neg_one(&self) -> Self::Element;
 11    /// Largest unsigned value that fits in the modulus. That is, q - 1.
 12    fn largest_unsigned_value(&self) -> Self::Element;
 13    /// Smallest unsigned value that fits in the modulus
 14    /// Always assmed to be 0.
 15    fn smallest_unsigned_value(&self) -> Self::Element;
 16    /// Max +value in signed representation
 17    fn signed_max(&self) -> Self::Element;
 18    /// Min -value in signed representation
 19    fn signed_min(&self) -> Self::Element;
 20}
 21
 22pub trait ModInit {
 23    type Element;
 24    fn new(q: Self::Element) -> Self;
 25}
 26
 27pub trait VectorOps {
 28    type Element;
 29
 30    fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element);
 31    fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]);
 32
 33    fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
 34    fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
 35    fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
 36    fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element);
 37    fn elwise_neg_mut(&self, a: &mut [Self::Element]);
 38    /// inplace mutates `a`: a = a + b*c
 39    fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]);
 40    fn elwise_fma_scalar_mut(
 41        &self,
 42        a: &mut [Self::Element],
 43        b: &[Self::Element],
 44        c: &Self::Element,
 45    );
 46
 47    fn modulus(&self) -> Self::Element;
 48}
 49
 50pub trait ArithmeticOps {
 51    type Element;
 52
 53    fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
 54    fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
 55    fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
 56    fn neg(&self, a: &Self::Element) -> Self::Element;
 57
 58    fn modulus(&self) -> Self::Element;
 59}
 60
 61pub struct ModularOpsU64 {
 62    q: u64,
 63    logq: usize,
 64    barrett_mu: u128,
 65    barrett_alpha: usize,
 66}
 67
 68impl ModInit for ModularOpsU64 {
 69    type Element = u64;
 70    fn new(q: u64) -> ModularOpsU64 {
 71        let logq = 64 - q.leading_zeros();
 72
 73        // barrett calculation
 74        let mu = (1u128 << (logq * 2 + 3)) / (q as u128);
 75        let alpha = logq + 3;
 76
 77        ModularOpsU64 {
 78            q,
 79            logq: logq as usize,
 80            barrett_alpha: alpha as usize,
 81            barrett_mu: mu,
 82        }
 83    }
 84}
 85
 86impl ModularOpsU64 {
 87    fn add_mod_fast(&self, a: u64, b: u64) -> u64 {
 88        debug_assert!(a < self.q);
 89        debug_assert!(b < self.q);
 90
 91        let mut o = a + b;
 92        if o >= self.q {
 93            o -= self.q;
 94        }
 95        o
 96    }
 97
 98    fn sub_mod_fast(&self, a: u64, b: u64) -> u64 {
 99        debug_assert!(a < self.q);
100        debug_assert!(b < self.q);
101
102        if a > b {
103            a - b
104        } else {
105            (self.q + a) - b
106        }
107    }
108
109    /// returns (a * b)  % q
110    ///
111    /// - both a and b must be in range [0, 2q)
112    /// - output is in range [0 , q)
113    fn mul_mod_fast(&self, a: u64, b: u64) -> u64 {
114        debug_assert!(a < 2 * self.q);
115        debug_assert!(b < 2 * self.q);
116
117        let ab = a as u128 * b as u128;
118
119        // ab / (2^{n + \beta})
120        // note: \beta is assumed to -2
121        let tmp = ab >> (self.logq - 2);
122
123        // k = ((ab / (2^{n + \beta})) * \mu) / 2^{\alpha - (-2)}
124        let k = (tmp * self.barrett_mu) >> (self.barrett_alpha + 2);
125
126        // ab - k*p
127        let tmp = k * (self.q as u128);
128
129        let mut out = (ab - tmp) as u64;
130
131        if out >= self.q {
132            out -= self.q;
133        }
134
135        return out;
136    }
137}
138
139impl ArithmeticOps for ModularOpsU64 {
140    type Element = u64;
141
142    fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
143        self.add_mod_fast(*a, *b)
144    }
145
146    fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
147        self.mul_mod_fast(*a, *b)
148    }
149
150    fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
151        self.sub_mod_fast(*a, *b)
152    }
153
154    fn neg(&self, a: &Self::Element) -> Self::Element {
155        self.q - *a
156    }
157
158    fn modulus(&self) -> Self::Element {
159        self.q
160    }
161}
162
163impl VectorOps for ModularOpsU64 {
164    type Element = u64;
165
166    fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
167        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
168            *ai = self.add_mod_fast(*ai, *bi);
169        });
170    }
171
172    fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
173        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
174            *ai = self.sub_mod_fast(*ai, *bi);
175        });
176    }
177
178    fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
179        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
180            *ai = self.mul_mod_fast(*ai, *bi);
181        });
182    }
183
184    fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
185        a.iter_mut().for_each(|ai| *ai = self.q - *ai);
186    }
187
188    fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
189        izip!(out.iter_mut(), a.iter()).for_each(|(oi, ai)| {
190            *oi = self.mul_mod_fast(*ai, *b);
191        });
192    }
193
194    fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]) {
195        izip!(out.iter_mut(), a.iter(), b.iter()).for_each(|(oi, ai, bi)| {
196            *oi = self.mul_mod_fast(*ai, *bi);
197        });
198    }
199
200    fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element) {
201        a.iter_mut().for_each(|ai| {
202            *ai = self.mul_mod_fast(*ai, *b);
203        });
204    }
205
206    fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
207        izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(ai, bi, ci)| {
208            *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *ci));
209        });
210    }
211
212    fn elwise_fma_scalar_mut(
213        &self,
214        a: &mut [Self::Element],
215        b: &[Self::Element],
216        c: &Self::Element,
217    ) {
218        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
219            *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *c));
220        });
221    }
222
223    fn modulus(&self) -> Self::Element {
224        self.q
225    }
226}
227
228pub struct WordSizeModulus<T> {
229    _phantom: PhantomData<T>,
230}
231
232impl<T> ModInit for WordSizeModulus<T> {
233    type Element = T;
234    fn new<M>(q: M) -> Self {
235        // For now assume ModulusOpsU64 is only used for u64
236        Self {
237            _phantom: PhantomData,
238        }
239    }
240}
241
242// impl<T: WrappingAdd + WrappingSub + WrappingMul> ArithmeticOps for
243// WordSizeModulus<T> {     fn add(&self, a: &Self::Element, b: &Self::Element)
244// -> Self::Element {         T::wrapping_add(*a, *b)
245//     }
246
247//     fn modulus(&self) -> Self::Element {
248
249//     }
250// }