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