1use crate::modulus::barrett::BarrettPrecomp;
  2use crate::modulus::montgomery::{Montgomery, MontgomeryPrecomp};
  3use crate::modulus::ReduceOnce;
  4use crate::modulus::{ONCE, REDUCEMOD};
  5extern crate test;
  6
  7/// MontgomeryPrecomp is a set of methods implemented for MontgomeryPrecomp<u64>
  8/// enabling Montgomery arithmetic over u64 values.
  9impl MontgomeryPrecomp<u64> {
 10    /// Returns an new instance of MontgomeryPrecomp<u64>.
 11    /// This method will fail if gcd(q, 2^64) != 1.
 12    #[inline(always)]
 13    pub fn new(q: u64) -> MontgomeryPrecomp<u64> {
 14        assert!(
 15            q & 1 != 0,
 16            "Invalid argument: gcd(q={}, radix=2^64) != 1",
 17            q
 18        );
 19        let mut q_inv: u64 = 1;
 20        let mut q_pow = q;
 21        for _i in 0..63 {
 22            q_inv = q_inv.wrapping_mul(q_pow);
 23            q_pow = q_pow.wrapping_mul(q_pow);
 24        }
 25        let mut precomp = Self {
 26            q: q,
 27            two_q: q << 1,
 28            four_q: q << 2,
 29            barrett: BarrettPrecomp::new(q),
 30            q_inv: q_inv,
 31            one: 0,
 32            minus_one: 0,
 33        };
 34
 35        precomp.one = precomp.prepare::<ONCE>(1);
 36        precomp.minus_one = q - precomp.one;
 37
 38        precomp
 39    }
 40
 41    /// Returns 2^64 mod q as a Montgomery<u64>.
 42    #[inline(always)]
 43    pub fn one(&self) -> Montgomery<u64> {
 44        self.one
 45    }
 46
 47    /// Returns (q-1) * 2^64 mod q as a Montgomery<u64>.
 48    #[inline(always)]
 49    pub fn minus_one(&self) -> Montgomery<u64> {
 50        self.minus_one
 51    }
 52
 53    /// Applies a modular reduction on x based on REDUCE:
 54    /// - LAZY: no modular reduction.
 55    /// - ONCE: subtracts q if x >= q.
 56    /// - FULL: maps x to x mod q using Barrett reduction.
 57    #[inline(always)]
 58    pub fn reduce<const REDUCE: REDUCEMOD>(&self, x: u64) -> u64 {
 59        let mut r: u64 = x;
 60        self.reduce_assign::<REDUCE>(&mut r);
 61        r
 62    }
 63
 64    /// Applies a modular reduction on x based on REDUCE:
 65    /// - LAZY: no modular reduction.
 66    /// - ONCE: subtracts q if x >= q.
 67    /// - FULL: maps x to x mod q using Barrett reduction.
 68    #[inline(always)]
 69    pub fn reduce_assign<const REDUCE: REDUCEMOD>(&self, x: &mut u64) {
 70        self.barrett.reduce_assign::<REDUCE>(x);
 71    }
 72
 73    /// Returns lhs * 2^64 mod q as a Montgomery<u64>.
 74    #[inline(always)]
 75    pub fn prepare<const REDUCE: REDUCEMOD>(&self, lhs: u64) -> Montgomery<u64> {
 76        let mut rhs: u64 = 0;
 77        self.prepare_assign::<REDUCE>(lhs, &mut rhs);
 78        rhs
 79    }
 80
 81    /// Assigns lhs * 2^64 mod q to rhs.
 82    #[inline(always)]
 83    pub fn prepare_assign<const REDUCE: REDUCEMOD>(&self, lhs: u64, rhs: &mut Montgomery<u64>) {
 84        let (_, mhi) = lhs.widening_mul(*self.barrett.value_lo());
 85        *rhs = (lhs.wrapping_mul(*self.barrett.value_hi()).wrapping_add(mhi))
 86            .wrapping_mul(self.q)
 87            .wrapping_neg();
 88        self.reduce_assign::<REDUCE>(rhs);
 89    }
 90
 91    /// Returns lhs * (2^64)^-1 mod q as a u64.
 92    #[inline(always)]
 93    pub fn unprepare<const REDUCE: REDUCEMOD>(&self, lhs: Montgomery<u64>) -> u64 {
 94        let mut rhs = 0u64;
 95        self.unprepare_assign::<REDUCE>(lhs, &mut rhs);
 96        rhs
 97    }
 98
 99    /// Assigns lhs * (2^64)^-1 mod q to rhs.
100    #[inline(always)]
101    pub fn unprepare_assign<const REDUCE: REDUCEMOD>(&self, lhs: Montgomery<u64>, rhs: &mut u64) {
102        let (_, r) = self.q.widening_mul(lhs.wrapping_mul(self.q_inv));
103        *rhs = self.reduce::<REDUCE>(self.q.wrapping_sub(r));
104    }
105
106    /// Returns lhs * rhs * (2^{64})^-1 mod q.
107    #[inline(always)]
108    pub fn mul_external<const REDUCE: REDUCEMOD>(&self, lhs: Montgomery<u64>, rhs: u64) -> u64 {
109        let mut r: u64 = rhs;
110        self.mul_external_assign::<REDUCE>(lhs, &mut r);
111        r
112    }
113
114    /// Assigns lhs * rhs * (2^{64})^-1 mod q to rhs.
115    #[inline(always)]
116    pub fn mul_external_assign<const REDUCE: REDUCEMOD>(
117        &self,
118        lhs: Montgomery<u64>,
119        rhs: &mut u64,
120    ) {
121        let (mlo, mhi) = lhs.widening_mul(*rhs);
122        let (_, hhi) = self.q.widening_mul(mlo.wrapping_mul(self.q_inv));
123        *rhs = self.reduce::<REDUCE>(mhi.wrapping_sub(hhi).wrapping_add(self.q));
124    }
125
126    /// Returns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1].
127    #[inline(always)]
128    pub fn mul_internal<const REDUCE: REDUCEMOD>(
129        &self,
130        lhs: Montgomery<u64>,
131        rhs: Montgomery<u64>,
132    ) -> Montgomery<u64> {
133        self.mul_external::<REDUCE>(lhs, rhs)
134    }
135
136    /// Assigns lhs * rhs * (2^{64})^-1 mod q to rhs.
137    #[inline(always)]
138    pub fn mul_internal_assign<const REDUCE: REDUCEMOD>(
139        &self,
140        lhs: Montgomery<u64>,
141        rhs: &mut Montgomery<u64>,
142    ) {
143        self.mul_external_assign::<REDUCE>(lhs, rhs);
144    }
145
146    #[inline(always)]
147    pub fn add_internal(&self, lhs: Montgomery<u64>, rhs: Montgomery<u64>) -> Montgomery<u64> {
148        rhs + lhs
149    }
150
151    /// Assigns lhs + rhs to rhs.
152    #[inline(always)]
153    pub fn add_internal_lazy_assign(&self, lhs: Montgomery<u64>, rhs: &mut Montgomery<u64>) {
154        *rhs += lhs
155    }
156
157    /// Assigns lhs + rhs - q if (lhs + rhs) >= q to rhs.
158    #[inline(always)]
159    pub fn add_internal_reduce_once_assign<const LAZY: bool>(
160        &self,
161        lhs: Montgomery<u64>,
162        rhs: &mut Montgomery<u64>,
163    ) {
164        self.add_internal_lazy_assign(lhs, rhs);
165        rhs.reduce_once_assign(self.q);
166    }
167
168    /// Returns (x^exponent) * 2^64 mod q.
169    #[inline(always)]
170    pub fn pow(&self, x: Montgomery<u64>, exponent: u64) -> Montgomery<u64> {
171        let mut y: Montgomery<u64> = self.one();
172        let mut x_mut: Montgomery<u64> = x;
173        let mut i: u64 = exponent;
174        while i > 0 {
175            if i & 1 == 1 {
176                self.mul_internal_assign::<ONCE>(x_mut, &mut y);
177            }
178            self.mul_internal_assign::<ONCE>(x_mut, &mut x_mut);
179            i >>= 1;
180        }
181
182        y.reduce_once_assign(self.q);
183        y
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190    use crate::modulus::montgomery;
191    use test::Bencher;
192
193    #[test]
194    fn test_mul_external() {
195        let q: u64 = 0x1fffffffffe00001;
196        let m_precomp = montgomery::MontgomeryPrecomp::new(q);
197        let x: u64 = 0x5f876e514845cc8b;
198        let y: u64 = 0xad726f98f24a761a;
199        let y_mont = m_precomp.prepare::<ONCE>(y);
200        assert!(
201            m_precomp.mul_external::<ONCE>(y_mont, x) == (x as u128 * y as u128 % q as u128) as u64
202        );
203    }
204
205    #[bench]
206    fn bench_mul_external(b: &mut Bencher) {
207        let q: u64 = 0x1fffffffffe00001;
208        let m_precomp = montgomery::MontgomeryPrecomp::new(q);
209        let mut x: u64 = 0x5f876e514845cc8b;
210        let y: u64 = 0xad726f98f24a761a;
211        let y_mont = m_precomp.prepare::<ONCE>(y);
212        b.iter(|| m_precomp.mul_external_assign::<ONCE>(y_mont, &mut x));
213    }
214}