1use crate::modulus::barrett::BarrettPrecomp;
  2use crate::modulus::ReduceOnce;
  3
  4extern crate test;
  5
  6/// Montgomery is a generic struct storing
  7/// an element in the Montgomery domain.
  8#[derive(Clone, Copy, Debug, PartialEq, Eq)]
  9pub struct Montgomery<O>(pub O);
 10
 11/// Implements helper methods on the struct Montgomery<O>.
 12impl<O> Montgomery<O>{
 13
 14    #[inline(always)]
 15    pub fn new(lhs: O) -> Self{
 16        Self(lhs)
 17    }
 18
 19    #[inline(always)]
 20    pub fn value(&self) -> &O{
 21        &self.0
 22    }
 23
 24    pub fn value_mut(&mut self) -> &mut O{
 25        &mut self.0
 26    }
 27}
 28
 29/// MontgomeryPrecomp is a generic struct storing 
 30/// precomputations for Montgomery arithmetic.
 31#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 32pub struct MontgomeryPrecomp<O>{
 33    q: O,
 34    barrett: BarrettPrecomp<O>,
 35    q_inv: O,
 36    one: Montgomery<O>,
 37    minus_one: Montgomery<O>,
 38}
 39
 40/// MontgomeryPrecomp is a set of methods implemented for MontgomeryPrecomp<u64>
 41/// enabling Montgomery arithmetic over u64 values.
 42impl MontgomeryPrecomp<u64>{
 43
 44    /// Returns an new instance of MontgomeryPrecomp<u64>.
 45    /// This method will fail if gcd(q, 2^64) != 1.
 46    #[inline(always)]
 47    pub fn new(q: u64) -> MontgomeryPrecomp<u64>{
 48        assert!(q & 1 != 0, "Invalid argument: gcd(q={}, radix=2^64) != 1", q);
 49        let mut q_inv: u64 = 1;
 50        let mut q_pow = q;
 51        for _i in 0..63{
 52            q_inv = q_inv.wrapping_mul(q_pow);
 53            q_pow = q_pow.wrapping_mul(q_pow);
 54        }
 55        let mut precomp = Self{ 
 56            q: q, 
 57            barrett: BarrettPrecomp::new(q), 
 58            q_inv: q_inv,
 59            one: Montgomery(0),
 60            minus_one: Montgomery(0),
 61        };
 62
 63        precomp.one = precomp.prepare(1);
 64        precomp.minus_one = Montgomery(q-precomp.one.value());
 65
 66        precomp
 67    }
 68
 69    /// Returns 2^64 mod q as a Montgomery<u64>.
 70    #[inline(always)]
 71    pub fn one(&self) -> Montgomery<u64>{
 72        self.one
 73    }
 74
 75    /// Returns (q-1) * 2^64 mod q as a Montgomery<u64>.
 76    #[inline(always)]
 77    pub fn minus_one(&self) -> Montgomery<u64>{
 78        self.minus_one
 79    }
 80
 81    /// Returns lhs * 2^64 mod q as a Montgomery<u64>.
 82    #[inline(always)]
 83    pub fn prepare(&self, lhs: u64) -> Montgomery<u64>{
 84        let mut rhs = Montgomery(0);
 85        self.prepare_assign(lhs, &mut rhs);
 86        rhs
 87    }
 88
 89    /// Assigns lhs * 2^64 mod q to rhs.
 90    #[inline(always)]
 91    pub fn prepare_assign(&self, lhs: u64, rhs: &mut Montgomery<u64>){
 92        self.prepare_lazy_assign(lhs, rhs);
 93        rhs.value_mut().reduce_once_assign(self.q);
 94    }
 95
 96    /// Returns lhs * 2^64 mod q in range [0, 2q-1] as a Montgomery<u64>.
 97    #[inline(always)]
 98    pub fn prepare_lazy(&self, lhs: u64) -> Montgomery<u64>{
 99        let mut rhs = Montgomery(0);
100        self.prepare_lazy_assign(lhs, &mut rhs);
101        rhs
102    }
103
104    /// Assigns lhs * 2^64 mod q in range [0, 2q-1] to rhs.
105    #[inline(always)]
106    pub fn prepare_lazy_assign(&self, lhs: u64, rhs: &mut Montgomery<u64>){
107        let (_, mhi) = lhs.widening_mul(*self.barrett.value_lo());
108        *rhs = Montgomery((lhs.wrapping_mul(*self.barrett.value_hi()).wrapping_add(mhi)).wrapping_mul(self.q).wrapping_neg());
109    }
110
111    /// Returns lhs * (2^64)^-1 mod q as a u64.
112    #[inline(always)]
113    pub fn unprepare(&self, lhs: Montgomery<u64>) -> u64{
114        let mut rhs = 0u64;
115        self.unprepare_assign(lhs, &mut rhs);
116        rhs
117    }
118
119    /// Assigns lhs * (2^64)^-1 mod q to rhs.
120    #[inline(always)]
121    pub fn unprepare_assign(&self, lhs: Montgomery<u64>, rhs: &mut u64){
122        self.unprepare_lazy_assign(lhs, rhs);
123        rhs.reduce_once_assign(self.q);
124        
125    }
126
127    /// Returns lhs * (2^64)^-1 mod q in range [0, 2q-1].
128    #[inline(always)]
129    pub fn unprepare_lazy(&self, lhs: Montgomery<u64>) -> u64{
130        let mut rhs = 0u64;
131        self.unprepare_lazy_assign(lhs, &mut rhs);
132        rhs
133    }
134
135    /// Assigns lhs * (2^64)^-1 mod q in range [0, 2q-1] to rhs.
136    #[inline(always)]
137    pub fn unprepare_lazy_assign(&self, lhs: Montgomery<u64>, rhs: &mut u64){
138        let (_, r) = self.q.widening_mul(lhs.value().wrapping_mul(self.q_inv));
139        *rhs = self.q - r
140    }
141
142    /// Returns lhs * rhs * (2^{64})^-1 mod q.
143    #[inline(always)]
144    pub fn mul_external(&self, lhs: Montgomery<u64>, rhs: u64) -> u64{
145        let mut r = self.mul_external_lazy(lhs, rhs);
146        r.reduce_once_assign(self.q);
147        r
148    }
149
150    /// Assigns lhs * rhs * (2^{64})^-1 mod q to rhs.
151    #[inline(always)]
152    pub fn mul_external_assign(&self, lhs: Montgomery<u64>, rhs: &mut u64){
153        self.mul_external_lazy_assign(lhs, rhs);
154        rhs.reduce_once_assign(self.q);
155    }
156
157    /// Returns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1].
158    #[inline(always)]
159    pub fn mul_external_lazy(&self, lhs: Montgomery<u64>, rhs: u64) -> u64{
160        let mut result = rhs;
161        self.mul_external_lazy_assign(lhs, &mut result);
162        result
163    }
164
165    /// Assigns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1] to rhs.
166    #[inline(always)]
167    pub fn mul_external_lazy_assign(&self, lhs: Montgomery<u64>, rhs: &mut u64){
168        let (mlo, mhi) = lhs.value().widening_mul(*rhs);
169        let (_, hhi) = self.q.widening_mul(mlo.wrapping_mul(self.q_inv));
170        *rhs = mhi.wrapping_add(self.q - hhi)
171    }
172
173    /// Returns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1].
174    #[inline(always)]
175    pub fn mul_internal(&self, lhs: Montgomery<u64>, rhs: Montgomery<u64>) -> Montgomery<u64>{
176        Montgomery(self.mul_external(lhs, *rhs.value()))
177    }
178
179    /// Assigns lhs * rhs * (2^{64})^-1 mod q to rhs.
180    #[inline(always)]
181    pub fn mul_internal_assign(&self, lhs: Montgomery<u64>, rhs: &mut Montgomery<u64>){
182        self.mul_external_assign(lhs, rhs.value_mut());
183    }
184
185    /// Returns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1].
186    #[inline(always)]
187    pub fn mul_internal_lazy(&self, lhs: Montgomery<u64>, rhs: Montgomery<u64>) -> Montgomery<u64>{
188        Montgomery(self.mul_external_lazy(lhs, *rhs.value()))
189    }
190
191    /// Assigns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1] to rhs.
192    #[inline(always)]
193    pub fn mul_internal_lazy_assign(&self, lhs: Montgomery<u64>, rhs: &mut Montgomery<u64>){
194        self.mul_external_lazy_assign(lhs, rhs.value_mut());
195    }
196
197    #[inline(always)]
198    pub fn add_internal(&self, lhs: Montgomery<u64>, rhs: Montgomery<u64>) -> Montgomery<u64>{
199        Montgomery(self.barrett.reduce(rhs.value() + lhs.value()))
200    }
201
202    /// Assigns lhs + rhs to rhs.
203    #[inline(always)]
204    pub fn add_internal_lazy_assign(&self, lhs: Montgomery<u64>, rhs: &mut Montgomery<u64>){
205        *rhs.value_mut() += lhs.value()
206    }
207
208    /// Assigns lhs + rhs - q if (lhs + rhs) >= q to rhs.
209    #[inline(always)]
210    pub fn add_internal_reduce_once_assign(&self, lhs: Montgomery<u64>, rhs: &mut Montgomery<u64>){
211        self.add_internal_lazy_assign(lhs, rhs);
212        rhs.value_mut().reduce_once_assign(self.q);
213    }
214
215    /// Returns (x^exponent) * 2^64 mod q.
216    #[inline(always)]
217    pub fn pow(&self, x: Montgomery<u64>, exponent:u64) -> Montgomery<u64>{
218        let mut y: Montgomery<u64> = self.one();
219        let mut x_mut: Montgomery<u64> = x;
220        let mut i: u64 = exponent;
221        while i > 0{
222            if i & 1 == 1{
223                self.mul_internal_assign(x_mut, &mut y);
224            }
225            self.mul_internal_assign(x_mut, &mut x_mut);
226            i >>= 1;
227        }
228
229        y.value_mut().reduce_once_assign(self.q);
230        y
231    }
232}
233
234/// Returns x^exponent mod q.
235/// This function internally instantiate a new MontgomeryPrecomp<u64>
236/// To be used when called only a few times and if there 
237/// is no Prime instantiated with q.
238fn pow(x:u64, exponent:u64, q:u64) -> u64{
239    let montgomery: MontgomeryPrecomp<u64> = MontgomeryPrecomp::<u64>::new(q);
240    let mut y_mont: Montgomery<u64> = montgomery.one();
241    let mut x_mont: Montgomery<u64> = montgomery.prepare(x);
242    while exponent > 0{
243        if exponent & 1 == 1{
244            montgomery.mul_internal_assign(x_mont, &mut y_mont);
245        }
246
247        montgomery.mul_internal_assign(x_mont, &mut x_mont);
248    }
249    
250    montgomery.unprepare(y_mont)
251}
252
253#[cfg(test)]
254mod tests {
255    use crate::modulus::montgomery;
256    use super::*;
257    use test::Bencher;
258
259    #[test]
260    fn test_mul_external() {
261        let q: u64 = 0x1fffffffffe00001;
262    	let m_precomp = montgomery::MontgomeryPrecomp::new(q);
263        let x: u64 = 0x5f876e514845cc8b;
264        let y: u64 = 0xad726f98f24a761a;
265        let y_mont = m_precomp.prepare(y);
266    	assert!(m_precomp.mul_external(y_mont, x) == (x as u128 * y as u128 % q as u128) as u64);
267    }
268
269    #[bench]
270    fn bench_mul_external(b: &mut Bencher){
271        let q: u64 = 0x1fffffffffe00001;
272    	let m_precomp = montgomery::MontgomeryPrecomp::new(q);
273        let mut x: u64 = 0x5f876e514845cc8b;
274        let y: u64 = 0xad726f98f24a761a;
275        let y_mont = m_precomp.prepare(y);
276        b.iter(|| m_precomp.mul_external_assign(y_mont, &mut x));
277    }
278}