From 1a7e57404ee0acf272b94147694862972ded9363 Mon Sep 17 00:00:00 2001 From: ddidderr Date: Fri, 9 Feb 2024 22:38:11 +0100 Subject: [PATCH] [code] small improvements * readability: changing variable names/adding variables for clarity * initialize primes Vec with reasonable space to reduce allocations * return type of find_next_prime more idiomatic (Option) * avoid overflow by using checked_mul --- src/sieve.rs | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/src/sieve.rs b/src/sieve.rs index 42038bc..eabe333 100644 --- a/src/sieve.rs +++ b/src/sieve.rs @@ -1,7 +1,7 @@ pub fn get_primes(max_nr: usize) -> Vec { let mut sieve = vec![true; max_nr + 1]; - let mut primes = Vec::new(); + let mut primes = Vec::with_capacity(1_000_000); sieve[0] = false; sieve[1] = false; @@ -12,7 +12,7 @@ pub fn get_primes(max_nr: usize) -> Vec { loop { check_off_multiples_of_nr(nr, &mut sieve); - if !find_next_prime(&mut nr, &mut sieve) { + if find_next_prime(&mut nr, &mut sieve).is_none() { break; } @@ -22,24 +22,29 @@ pub fn get_primes(max_nr: usize) -> Vec { primes } -fn find_next_prime(nr: &mut u64, sieve: &mut [bool]) -> bool { +fn find_next_prime(nr: &mut u64, sieve: &mut [bool]) -> Option { + let start_idx = usize::try_from(*nr).ok()? + 1; + if let Some((idx, _)) = sieve .iter() .enumerate() - .skip(usize::try_from(*nr).unwrap() + 1) - .find(|(_, val)| **val) + .skip(start_idx) + .find(|&(_, &is_prime)| is_prime) { *nr = idx as u64; - return true; + return Some(*nr); } - false + None } fn check_off_multiples_of_nr(nr: u64, sieve: &mut [bool]) { + + let nr = usize::try_from(nr).unwrap(); + sieve .iter_mut() - .skip(usize::try_from(nr * 2).unwrap()) - .step_by(usize::try_from(nr).unwrap()) + .skip(nr.checked_mul(2).unwrap()) + .step_by(nr) .for_each(|nr| *nr = false); }