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