1use crate::modulus::barrett::BarrettPrecomp;
  2use crate::modulus::montgomery::{Montgomery, MontgomeryPrecomp};
  3use crate::modulus::prime::Prime;
  4use crate::modulus::ONCE;
  5use primality_test::is_prime;
  6use prime_factorization::Factorization;
  7
  8impl Prime<u64> {
  9    /// Returns a new instance of Prime<u64>.
 10    /// Panics if q_base is not a prime > 2 and
 11    /// if q_base^q_power would overflow u64.
 12    pub fn new(q_base: u64, q_power: usize) -> Self {
 13        assert!(is_prime(q_base) && q_base > 2);
 14        Self::new_unchecked(q_base, q_power)
 15    }
 16
 17    /// Returns a new instance of Prime<u64>.
 18    /// Does not check if q_base is a prime > 2.
 19    /// Panics if q_base^q_power would overflow u64.
 20    pub fn new_unchecked(q_base: u64, q_power: usize) -> Self {
 21        let mut q = q_base;
 22        for _i in 1..q_power {
 23            q *= q_base
 24        }
 25
 26        assert!(q.next_power_of_two().ilog2() <= 61);
 27
 28        let mut phi = q_base - 1;
 29        for _i in 1..q_power {
 30            phi *= q_base
 31        }
 32
 33        let mut prime: Prime<u64> = Self {
 34            q: q,
 35            two_q: q << 1,
 36            four_q: q << 2,
 37            q_base: q_base,
 38            q_power: q_power,
 39            factors: Vec::new(),
 40            montgomery: MontgomeryPrecomp::new(q),
 41            barrett: BarrettPrecomp::new(q),
 42            phi: phi,
 43        };
 44
 45        prime.check_factors();
 46
 47        prime
 48    }
 49
 50    pub fn q(&self) -> u64 {
 51        self.q
 52    }
 53
 54    pub fn q_base(&self) -> u64 {
 55        self.q_base
 56    }
 57
 58    pub fn q_power(&self) -> usize {
 59        self.q_power
 60    }
 61
 62    /// Returns x^exponen mod q.
 63    #[inline(always)]
 64    pub fn pow(&self, x: u64, exponent: u64) -> u64 {
 65        let mut y_mont: Montgomery<u64> = self.montgomery.one();
 66        let mut x_mont: Montgomery<u64> = self.montgomery.prepare::<ONCE>(x);
 67        let mut i: u64 = exponent;
 68        while i > 0 {
 69            if i & 1 == 1 {
 70                self.montgomery
 71                    .mul_internal_assign::<ONCE>(x_mont, &mut y_mont);
 72            }
 73
 74            self.montgomery
 75                .mul_internal_assign::<ONCE>(x_mont, &mut x_mont);
 76
 77            i >>= 1;
 78        }
 79
 80        self.montgomery.unprepare::<ONCE>(y_mont)
 81    }
 82
 83    /// Returns x^-1 mod q.
 84    /// User must ensure that x is not divisible by q_base.
 85    #[inline(always)]
 86    pub fn inv(&self, x: u64) -> u64 {
 87        self.pow(x, self.phi - 1)
 88    }
 89}
 90
 91impl Prime<u64> {
 92    /// Returns the smallest nth primitive root of q_base.
 93    pub fn primitive_root(&self) -> u64 {
 94        let mut candidate: u64 = 1u64;
 95        let mut not_found: bool = true;
 96
 97        while not_found {
 98            candidate += 1;
 99
100            for &factor in &self.factors {
101                if pow(candidate, (self.q_base - 1) / factor, self.q_base) == 1 {
102                    not_found = true;
103                    break;
104                }
105                not_found = false;
106            }
107        }
108
109        if not_found {
110            panic!("failed to find a primitive root for q_base={}", self.q_base)
111        }
112
113        candidate
114    }
115
116    /// Returns an nth primitive root of q = q_base^q_power in Montgomery.
117    pub fn primitive_nth_root(&self, nth_root: u64) -> u64 {
118        assert!(
119            self.q & (nth_root - 1) == 1,
120            "invalid prime: q = {} % nth_root = {} = {} != 1",
121            self.q,
122            nth_root,
123            self.q & (nth_root - 1)
124        );
125
126        let psi: u64 = self.primitive_root();
127
128        // nth primitive root mod q_base: psi_nth^(prime.q_base-1)/nth_root mod q_base
129        let psi_nth_q_base: u64 = pow(psi, (self.q_base - 1) / nth_root, self.q_base);
130
131        // lifts nth primitive root mod q_base to q = q_base^q_power
132        let psi_nth_q: u64 = self.hensel_lift(psi_nth_q_base, nth_root);
133
134        assert!(
135            self.pow(psi_nth_q, nth_root) == 1,
136            "invalid nth primitive root: psi^nth_root != 1 mod q"
137        );
138        assert!(
139            self.pow(psi_nth_q, nth_root >> 1) == self.q - 1,
140            "invalid nth primitive root: psi^(nth_root/2) != -1 mod q"
141        );
142
143        psi_nth_q
144    }
145
146    /// Checks if the field self.factor is populated.
147    /// If not, factorize q_base-1 and populates self.factor.
148    /// If yes, checks that it contains the unique factors of q_base-1.
149    pub fn check_factors(&mut self) {
150        if self.factors.len() == 0 {
151            let factors = Factorization::run(self.q_base - 1).prime_factor_repr();
152            let mut distincts_factors: Vec<u64> = Vec::with_capacity(factors.len());
153            for factor in factors.iter() {
154                distincts_factors.push(factor.0)
155            }
156            self.factors = distincts_factors
157        } else {
158            let mut q_base: u64 = self.q_base;
159
160            for &factor in &self.factors {
161                if !is_prime(factor) {
162                    panic!("invalid factor list: factor {} is not prime", factor)
163                }
164
165                while q_base % factor != 0 {
166                    q_base /= factor
167                }
168            }
169
170            if q_base != 1 {
171                panic!("invalid factor list: does not fully divide q_base: q_base % (all factors) = {}", q_base)
172            }
173        }
174    }
175
176    /// Returns (psi + a * q_base)^{nth_root} = 1 mod q = q_base^q_power given psi^{nth_root} = 1 mod q_base.
177    /// Panics if psi^{nth_root} != 1 mod q_base.
178    fn hensel_lift(&self, psi: u64, nth_root: u64) -> u64 {
179        assert!(
180            pow(psi, nth_root, self.q_base) == 1,
181            "invalid argument psi: psi^nth_root = {} != 1",
182            pow(psi, nth_root, self.q_base)
183        );
184
185        let mut psi_mont: Montgomery<u64> = self.montgomery.prepare::<ONCE>(psi);
186        let nth_root_mont: Montgomery<u64> = self.montgomery.prepare::<ONCE>(nth_root);
187
188        for _i in 1..self.q_power {
189            let psi_pow: Montgomery<u64> = self.montgomery.pow(psi_mont, nth_root - 1);
190
191            let num: Montgomery<u64> = self.montgomery.one() + self.q
192                - self.montgomery.mul_internal::<ONCE>(psi_pow, psi_mont);
193
194            let mut den: Montgomery<u64> =
195                self.montgomery.mul_internal::<ONCE>(nth_root_mont, psi_pow);
196
197            den = self.montgomery.pow(den, self.phi - 1);
198
199            psi_mont = self
200                .montgomery
201                .add_internal(psi_mont, self.montgomery.mul_internal::<ONCE>(num, den));
202        }
203
204        self.montgomery.unprepare::<ONCE>(psi_mont)
205    }
206}
207
208/// Returns x^exponent mod q.
209/// This function internally instantiate a new MontgomeryPrecomp<u64>
210/// To be used when called only a few times and if there
211/// is no Prime instantiated with q.
212pub fn pow(x: u64, exponent: u64, q: u64) -> u64 {
213    let montgomery: MontgomeryPrecomp<u64> = MontgomeryPrecomp::<u64>::new(q);
214    let mut y_mont: Montgomery<u64> = montgomery.one();
215    let mut x_mont: Montgomery<u64> = montgomery.prepare::<ONCE>(x);
216    let mut i: u64 = exponent;
217    while i > 0 {
218        if i & 1 == 1 {
219            montgomery.mul_internal_assign::<ONCE>(x_mont, &mut y_mont);
220        }
221
222        montgomery.mul_internal_assign::<ONCE>(x_mont, &mut x_mont);
223
224        i >>= 1;
225    }
226
227    montgomery.unprepare::<ONCE>(y_mont)
228}