Refractored code to handle import files better.

src/main.rs

@@ -22,24 +22,29 @@ use std::fs::File;
 use std::io::{BufRead, BufReader};
 use std::collections::VecDeque;
 
+mod test;
+
 #[derive(StructOpt)]
 #[structopt(about = "A prime number generator and tester.")]
-struct Opt {
+struct Opt
+{
 	#[structopt(short, long, help = "Print all found primes")]
 	verbose:bool,
 	#[structopt(short, long, name="FILE", help = "Import prime numbers from FILE")]
 	import:Option<PathBuf>,
+	#[structopt(short, long, help = "Test if n is prime instead of generation")]
+	test:bool,
 	#[structopt(help = "Ordinal of the prime to generate")]
 	n:usize,
 }
 
-fn main() {
+fn main()
+{
 	let opts = Opt::from_args();
 
 	// get the first `n` primes
 	let n = opts.n;
-	let mut primes:VecDeque<u64> = VecDeque::with_capacity(n);
-	let mut candidate:u64;
+	let mut primes:VecDeque<u64> = VecDeque::new();
 
 	if opts.import.is_some()
 	{
@@ -51,41 +56,87 @@ fn main() {
 			let aux:u64 = line.parse().unwrap();
 			primes.push_back(aux);
 		}
+	}
 
-		candidate = *primes.back().unwrap() + 2;
+	if opts.test
+	{
+		let mut res:bool;
+		if primes.len() == 0
+		{
+			res = test::is_prime(n as u64);
+		}
+		else if primes.back().unwrap() >= &(n as u64)
+		{
+			res = primes.contains(&(n as u64));
+		}
+		else if primes.back().unwrap() >= &((n as f64).sqrt() as u64)
+		{
+			res = test::is_prime_mem(n as u64, &primes)
 		}
 		else
 		{
-		// first prime
+			res = test::is_prime_mem(n as u64, &primes);
+			if res
+			{
+				res = test::is_prime_f(n as u64, primes.back().unwrap() + 2);
+			}
+		}
+		if res
+		{
+			if opts.verbose
+			{
+				println!("true");
+			}
+			std::process::exit(0);
+		}
+		else
+		{
+			if opts.verbose
+			{
+				println!("false");
+			}
+			std::process::exit(1);
+		}
+	}
+	else
+	{
+		if primes.len() >= n
+		{
+			println!("{}", primes.get(n-1).unwrap());
+		}
+		else
+		{
+			let mut candidate:u64;
+
+			if primes.len() == 0
+			{
+				primes.push_back(2);
 				if opts.verbose
 				{
 					println!("{}", 2);
 				}
-		primes.push_back(2);
 				candidate = 3;
 			}
+			else if primes.len() == 1
+			{
+				candidate = 3;
+			}
+			else
+			{
+				candidate = *primes.back().unwrap() + 2;
+			}
 
 			while primes.len() < n
 			{
-		let mut is_prime = true;
-		let limit = (candidate as f64).sqrt() as u64;
-		for i in primes.iter().take_while(|x| **x <= limit)
-		{
-			if candidate % *i == 0
-			{
-				is_prime = false;
-				break;
-			}
-		}
-
-		if is_prime
+				if test::is_prime_mem(candidate, &primes)
 				{
+					primes.push_back(candidate);
 					if opts.verbose
 					{
 						println!("{}", candidate);
 					}
-			primes.push_back(candidate);
 				}
+
 				candidate += 2;
 			}
 
@@ -94,3 +145,5 @@ fn main() {
 				println!("{}", primes.get(n-1).unwrap());
 			}
 		}
+	}
+}

src/test.rs

@@ -0,0 +1,70 @@
+/*
+ * Copyright (C) 2021  Ortega Froysa, Nicolás <nicolas@ortegas.org>
+ * Author: Ortega Froysa, Nicolás <nicolas@ortegas.org>
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+use std::collections::VecDeque;
+
+pub fn is_prime_f(n:u64, b:u64) -> bool
+{
+	assert_ne!(b, 0);
+	assert_ne!(b, 1);
+
+	let mut start = b;
+	if start == 2
+	{
+		if n % 2 == 0
+		{
+			return false;
+		}
+		else
+		{
+			start += 1;
+		}
+	}
+	else if start % 2 == 0
+	{
+		start += 1;
+	}
+
+	let limit = (n as f64).sqrt() as u64 + 1;
+	for i in (start..limit).step_by(2)
+	{
+		if n % i == 0
+		{
+			return false;
+		}
+	}
+	return true;
+}
+
+pub fn is_prime(n:u64) -> bool
+{
+	return is_prime_f(n, 2);
+}
+
+pub fn is_prime_mem(n:u64, primes:&VecDeque<u64>) -> bool
+{
+	let limit = (n as f64).sqrt() as u64;
+	for i in primes.iter().take_while(|x| **x <= limit)
+	{
+		if n % *i == 0
+		{
+			return false;
+		}
+	}
+	return true;
+}