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