1use crate::modulus::prime;
  2
  3use prime::Prime;
  4use primality_test::is_prime;
  5
  6impl NTTFriendlyPrimesGenerator<u64> {
  7	
  8	pub fn new(bit_size: u64, nth_root: u64) -> Self{
  9		let mut check_next_prime: bool = true;
 10		let mut check_prev_prime: bool = true;
 11		let next_prime = (1<<bit_size) + 1;
 12		let mut prev_prime = next_prime;
 13
 14		if next_prime > nth_root.wrapping_neg(){
 15			check_next_prime = false;
 16		}
 17
 18		if prev_prime < nth_root{
 19			check_prev_prime = false
 20		}
 21
 22		prev_prime -= nth_root;
 23
 24		Self{
 25			size: bit_size as f64,
 26			check_next_prime,
 27			check_prev_prime,
 28			nth_root,
 29			next_prime,
 30			prev_prime,
 31		}
 32	}
 33
 34	pub fn next_upstream_primes(&mut self, k: usize) -> Vec<Prime>{
 35		let mut primes: Vec<Prime> = Vec::with_capacity(k);
 36		for i in 0..k{
 37			primes.push(self.next_upstream_prime())
 38		}
 39		primes
 40	}
 41
 42	pub fn next_downstream_primes(&mut self, k: usize) -> Vec<Prime>{
 43		let mut primes: Vec<Prime> = Vec::with_capacity(k);
 44		for i in 0..k{
 45			primes.push(self.next_downstream_prime())
 46		}
 47		primes
 48	}
 49
 50	pub fn next_alternating_primes(&mut self, k: usize) -> Vec<Prime>{
 51		let mut primes: Vec<Prime> = Vec::with_capacity(k);
 52		for i in 0..k{
 53			primes.push(self.next_alternating_prime())
 54		}
 55		primes
 56	}
 57
 58	pub fn next_upstream_prime(&mut self) -> Prime{
 59		loop  {
 60			if self.check_next_prime{
 61				if (self.next_prime as f64).log2() - self.size >= 0.5 || self.next_prime > 0xffff_ffff_ffff_ffff-self.nth_root{
 62					self.check_next_prime = false;
 63					panic!("prime list for upstream primes is exhausted (overlap with next bit-size or prime > 2^64)");
 64				}
 65			}else{
 66				if is_prime(self.next_prime) {
 67					let prime = Prime::new_unchecked(self.next_prime);
 68					self.next_prime += self.nth_root;
 69					return prime
 70				}
 71				self.next_prime += self.nth_root;
 72			}
 73		}
 74	}
 75
 76	pub fn next_downstream_prime(&mut self) -> Prime{
 77		loop  {
 78			if self.size - (self.prev_prime as f64).log2() >= 0.5 || self.prev_prime < self.nth_root{
 79				self.check_next_prime = false;
 80				panic!("prime list for downstream primes is exhausted (overlap with previous bit-size or prime < nth_root)")
 81			}else{
 82				if is_prime(self.prev_prime){
 83					let prime = Prime::new_unchecked(self.next_prime);
 84					self.prev_prime -= self.nth_root;
 85					return prime
 86				}
 87				self.prev_prime -= self.nth_root; 
 88			}
 89		}
 90	}
 91
 92	pub fn next_alternating_prime(&mut self) -> Prime{
 93		loop {
 94			if !(self.check_next_prime || self.check_prev_prime){
 95				panic!("prime list for upstream and downstream prime is exhausted for the (overlap with previous/next bit-size or NthRoot > prime > 2^64)")
 96			}
 97
 98			if self.check_next_prime{
 99				if (self.next_prime as f64).log2() - self.size >= 0.5 || self.next_prime > 0xffff_ffff_ffff_ffff-self.nth_root{
100					self.check_next_prime = false;
101				}else{
102					if is_prime(self.next_prime){
103						let prime = Prime::new_unchecked(self.next_prime);
104						self.next_prime += self.nth_root;
105						return prime
106					}
107					self.next_prime += self.nth_root;
108				}
109			}
110
111			if self.check_prev_prime {
112				if self.size - (self.prev_prime as f64).log2() >= 0.5 || self.prev_prime < self.nth_root{
113					self.check_prev_prime = false;
114				}else{
115					if is_prime(self.prev_prime){
116						let prime = Prime::new_unchecked(self.prev_prime);
117						self.prev_prime -= self.nth_root;
118						return prime
119					}
120					self.prev_prime -= self.nth_root;
121				}
122			}
123		}
124	}
125}
126
127
128#[cfg(test)]
129mod test {
130    use crate::modulus::prime_generator;
131
132    #[test]
133    fn prime_generation() {
134    	let nth_root: u64 = 1<<16 ;
135    	let mut g: prime_generator::NTTFriendlyPrimesGenerator = prime_generator::NTTFriendlyPrimesGenerator::new(30, nth_root);
136
137    	let primes = g.next_alternating_primes(10);
138    	println!("{:?}", primes);
139    	for prime in primes.iter(){
140    		assert!(prime.q() % nth_root == 1);
141    	}
142    }
143}