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