1use crate::modulus::barrett::BarrettPrecomp;
  2use crate::modulus::ReduceOnce;
  3
  4/// Montgomery is a generic struct storing
  5/// an element in the Montgomery domain.
  6pub struct Montgomery<O>(O);
  7
  8/// Implements helper methods on the struct Montgomery<O>.
  9impl<O> Montgomery<O>{
 10
 11    #[inline(always)]
 12    pub fn new(lhs: O) -> Self{
 13        Self(lhs)
 14    }
 15
 16    #[inline(always)]
 17    pub fn value(&self) -> &O{
 18        &self.0
 19    }
 20
 21    pub fn value_mut(&mut self) -> &mut O{
 22        &mut self.0
 23    }
 24}
 25
 26/// MontgomeryPrecomp is a generic struct storing 
 27/// precomputations for Montgomery arithmetic.
 28pub struct MontgomeryPrecomp<O>{
 29    q: O,
 30    q_barrett: BarrettPrecomp<O>,
 31    q_inv: O,
 32}
 33
 34/// MontgomeryPrecomp is a set of methods implemented for MontgomeryPrecomp<u64>
 35/// enabling Montgomery arithmetic over u64 values.
 36impl MontgomeryPrecomp<u64>{
 37
 38    /// Returns an new instance of MontgomeryPrecomp<u64>.
 39    /// This method will fail if gcd(q, 2^64) != 1.
 40    #[inline(always)]
 41    fn new(&self, q: u64) -> MontgomeryPrecomp<u64>{
 42        assert!(q & 1 != 0, "Invalid argument: gcd(q={}, radix=2^64) != 1", q);
 43        let mut r: u64 = 1;
 44        let mut q_pow = q;
 45        for _i in 0..63{
 46            r = r.wrapping_mul(r);
 47            q_pow = q_pow.wrapping_mul(q_pow)
 48        }
 49        Self{ q: q, q_barrett: BarrettPrecomp::new(q, q), q_inv: q_pow}
 50    }
 51
 52    /// Returns (lhs<<64)%q as a Montgomery<u64>.
 53    #[inline(always)]
 54    fn prepare(&self, lhs: u64) -> Montgomery<u64>{
 55        let mut rhs = Montgomery(0);
 56        self.prepare_assign(lhs, &mut rhs);
 57        rhs
 58    }
 59
 60    /// Assigns (lhs<<64)%q to rhs.
 61    #[inline(always)]
 62    fn prepare_assign(&self, lhs: u64, rhs: &mut Montgomery<u64>){
 63        self.prepare_lazy_assign(lhs, rhs);
 64        rhs.value_mut().reduce_once_assign(self.q);
 65    }
 66
 67    /// Returns (lhs<<64)%q in range [0, 2q-1] as a Montgomery<u64>.
 68    #[inline(always)]
 69    fn prepare_lazy(&self, lhs: u64) -> Montgomery<u64>{
 70        let mut rhs = Montgomery(0);
 71        self.prepare_lazy_assign(lhs, &mut rhs);
 72        rhs
 73    }
 74
 75    /// Assigns (lhs<<64)%q in range [0, 2q-1] to rhs.
 76    #[inline(always)]
 77    fn prepare_lazy_assign(&self, lhs: u64, rhs: &mut Montgomery<u64>){
 78        let (mhi, _) = lhs.widening_mul(*self.q_barrett.value_lo());
 79        *rhs = Montgomery((lhs.wrapping_mul(*self.q_barrett.value_hi()).wrapping_add(mhi)).wrapping_mul(self.q).wrapping_neg());
 80    }
 81
 82    /// Returns lhs * rhs * (2^{64})^-1 mod q.
 83    #[inline(always)]
 84    fn mul_external(&self, lhs: Montgomery<u64>, rhs: u64) -> u64{
 85        let mut r = self.mul_external_lazy(lhs, rhs);
 86        r.reduce_once_assign(self.q);
 87        r
 88    }
 89
 90    /// Assigns lhs * rhs * (2^{64})^-1 mod q to rhs.
 91    #[inline(always)]
 92    fn mul_external_assign(&self, lhs: Montgomery<u64>, rhs: &mut u64){
 93        self.mul_external_lazy_assign(lhs, rhs);
 94        rhs.reduce_once_assign(self.q);
 95    }
 96
 97    /// Returns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1].
 98    #[inline(always)]
 99    fn mul_external_lazy(&self, lhs: Montgomery<u64>, rhs: u64) -> u64{
100        let mut result = rhs;
101        self.mul_external_lazy_assign(lhs, &mut result);
102        result
103    }
104
105    /// Assigns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1] to rhs.
106    #[inline(always)]
107    fn mul_external_lazy_assign(&self, lhs: Montgomery<u64>, rhs: &mut u64){
108        let (mhi, mlo) = lhs.value().widening_mul(*rhs);
109        let (hhi, _) = self.q.widening_mul(mlo * self.q_inv);
110        *rhs = mhi - hhi + self.q
111    }
112
113    /// Returns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1].
114    #[inline(always)]
115    fn mul_internal(&self, lhs: Montgomery<u64>, rhs: Montgomery<u64>) -> Montgomery<u64>{
116        Montgomery(self.mul_external(lhs, *rhs.value()))
117    }
118
119    /// Assigns lhs * rhs * (2^{64})^-1 mod q to rhs.
120    #[inline(always)]
121    fn mul_internal_assign(&self, lhs: Montgomery<u64>, rhs: &mut Montgomery<u64>){
122        self.mul_external_assign(lhs, rhs.value_mut());
123    }
124
125    /// Returns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1].
126    #[inline(always)]
127    fn mul_internal_lazy(&self, lhs: Montgomery<u64>, rhs: Montgomery<u64>) -> Montgomery<u64>{
128        Montgomery(self.mul_external_lazy(lhs, *rhs.value()))
129    }
130
131    /// Assigns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1] to rhs.
132    #[inline(always)]
133    fn mul_internal_lazy_assign(&self, lhs: Montgomery<u64>, rhs: &mut Montgomery<u64>){
134        self.mul_external_lazy_assign(lhs, rhs.value_mut());
135    }
136}