Improving the Prime Iterator in Rust

I wanted to improve the performance of the prime iterator I created back in 2021. The final code from part 1 was:

pub fn is_prime(n: u64) -> bool {
  if n < 4 {
    n > 1
  } else if n % 2 == 0 || n % 3 == 0 {
    false
  } else {
    let max_p = (n as f64).sqrt().ceil() as u64;
    match (5..=max_p).step_by(6).find(|p| n % p == 0 || n % (p+2) == 0) {
      Some(_) => false,
      None => true
    }
  }
}

pub struct Prime {
  curr: u64,
  next: u64,
  trial1: u64,
  trial2: u64
}

impl Prime {
  pub fn new() -> Prime {
    Prime {
      curr: 2,
      next: 3,
      trial1: 5,
      trial2: 7
    }
  }
}

impl Iterator for Prime {
  type Item = u64;

  fn next(&mut self) -> Option<Self::Item> {
    let prime = self.curr;
    self.curr = self.next;
    loop {
      self.next = self.trial1;
      self.trial1 = self.trial2;
      self.trial2 = self.next+6;
      if is_prime(self.next) {
        break;
      }
    }
    Some(prime)
  }
}

This algorithm operates in constant memory and uses the 6k ± 1 rule to generate the set of prime candidates. The area I would like to focus on is the is_prime function. is_prime needs to regenerate the list of primes each time to find out if the next prime candidate is prime or not. It will do this for each prime candidate tested. It would be faster if we memoize the generated primes so that we do not generate them over and over again. This will use more memory, but will perform faster.

Let's look at a version of the Iterator with memoization.

pub struct PrimeM {
  curr: u64,
  next: u64,
  trial1: u64,
  trial2: u64,
  primes: LinkedList<u64>
}

impl PrimeM {
  pub fn new() -> PrimeM {
    let mut prime_list = LinkedList::new();
    prime_list.push_back(2);
    prime_list.push_back(3);
    PrimeM {
      curr: 2,
      next: 3,
      trial1: 5,
      trial2: 7,
      primes: prime_list
    }
  }

  fn check_prime(&self, candidate:u64 ) -> bool {
    for p in &self.primes {
      let prime = *p;
      if prime*prime > candidate {
        return true;
      }
      if candidate%prime == 0 {
        return false;
      }
    }
    return true;
  }
}

impl Iterator for PrimeM {
  type Item = u64;

  fn next(&mut self) -> Option<Self::Item> {
    let prime = self.curr;
    self.curr = self.next;
    loop {
      self.next = self.trial1;
      self.trial1 = self.trial2;
      self.trial2 = self.next+6;
      if self.check_prime(self.next) {
        self.primes.push_back(self.next);
        break;
      }
    }
    Some(prime)
  }
}

There is now a linked list to store the generated primes. It is initialized with the first two primes 2 and 3. Each time the next prime is found it is stored in the linked list. A check_prime function has been added that loops through all the generated primes up to the square root of the candidate to see if the candidate is evenly divisible by the already generated primes. If it is not divisible then it is a prime number.

Here are the results of the two iterators generating 100000 primes:

The memoized version improved the performance of the iterator and runs in 29.47% of the time of the non-memoized version.

Full code can be found on GitHub

Blog

Improving the Prime Iterator in Rust

David Berry

Join Our Newsletter. No Spam, Only the good stuff.

Related