3
\$\begingroup\$

I wrote a prime number sieve in Ruby last year as part of a coding challenge, but lately I wanted to port it to Rust. I finally got around to it, but it's still intensely Rubinic.

Before I take the obvious next step and parallelize it, I want to make sure my code is as Rust-idiomatic as possible, so that I don't develop bad Rust habits.

Docs and tests can be run with Cargo from the source repository.

Code

It consists of three parts: a SquareMultiple iterator that computes the sequence \$i^2, i^2+i, i^2 + 2i...\$, a sieve of Eratosthenes, and then a segmented sieve that uses multiples from the Eratosthenes sieve to sieve the rest of the numbers up to the first function argument.

SquareMultiple

The iterator code came from an answer to this SO question.

/// The `prime-suspects` crate provides functions that sieve primes.
extern crate itertools;

use itertools::Itertools;

/// A struct that is used to generate the sequence j = i^2, i^2+i, i^2+2i,
/// i^2+3i, ..., used in the sieve of Eratosthenes.
struct SquareMultiple {
  /// The current iterator value. Starts at `arg * arg`, where `arg` is the
  /// single argument given to the constructor.
  curr: usize,
  /// The increment is just `arg`. It is added each time the iterator's .next()
  /// method is called.
  inc: usize
}

/// An `Iterator` implementation of `SquareMultiple`.
impl Iterator for SquareMultiple {
  type Item = usize;

  /// The return type is `Option<T>`:
  ///   * When the `Iterator` is finished, `None` is returned.
  ///   * Otherwise, the next value is wrapped in `Some` and returned.
  /// 
  /// Since there's no endpoint to this sequence, this specific `Iterator` will
  /// never return `None`. `Some` is always returned.
  ///
  /// The iterator uses only `.curr` and `.inc`, because all there is to do is
  /// add another of the original value (initialized in the implementation).
  fn next(&mut self) -> Option<usize> {
    let val = self.curr;
    self.curr += self.inc;
    Some(val)
  }
}

/// The SquareMultiple implementation itself, using the struct and Iterator
/// traits defined above.
impl SquareMultiple {
  fn new(term: usize) -> Self {
    SquareMultiple { curr: term * term, inc: term }
  }
}

Eratosthenes' Sieve

This code was a near-straight implementation of the Ruby version, which follows the pseudocode in the Wikipedia article.

/// An implementation of the sieve of Eratosthenes, as described in [the
/// Wikipedia
/// article](https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes#Algorithm_and_variants).
///
/// # Examples
///  ```
///  assert_eq!(vec![2,3,5,7], prime_suspects::eratosthenes_sieve(10));
///  ```
///
/// ```
/// assert_eq!(vec![65521, 65519, 65497],
/// prime_suspects::eratosthenes_sieve(65535)
///   .iter().rev().take(3).map(|&num| num)
///   .collect::<Vec<usize>>());
/// ```
pub fn eratosthenes_sieve(max_val: usize) -> Vec<usize> {
  // Algorithm notes: The sieve works like so (Wikipedia, pseudocode inlined):
  // Input: an integer n > 1
  // Let A be an array of Boolean values, indexed by integers 2 to n,
  // initially all set to true.
  let mut bool_vec = vec![true; max_val];
  // for i = 2, 3, 4, ..., not exceeding √n:
  let mut top_sieve = max_val as f64;
  // We have to add 1 because the sqrt coerced to an int is √floor(n)
  top_sieve = top_sieve.sqrt() + 1.0;
  for sieve_term in 2..(top_sieve as usize) {
    // if A[i] is true:
    if bool_vec[sieve_term] == true {
      // for j = i^2, i^2+i, i^2+2i, i^2+3i, ...,
      for j in SquareMultiple::new(sieve_term)
        .take_while(|&term| term < max_val) { // ...not exceeding n
          bool_vec[j] = false; // A[j] := false
        }
    }
  }
  let mut out_vec = vec![];
  // Output: all i such that A[i] is true.
  for term in 2..max_val {
    if bool_vec[term] == true {
      out_vec.push(term);
    }
  }
  out_vec
}

Segmented Sieve

I struggled with this one the most. I wanted a .each_slice() method very badly, since that was the quickest route to the same algorithm in Ruby, and asked for it on SO.

/// A [segmented
/// approach](https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes#Segmented_sieve)
/// to sieveing, keeping memory use to O(√n). As Sorensen states, this is the
/// most practical optimization to the sieve of Eratosthenes.
///
/// # Examples
///
/// ```
/// assert_eq!(vec![65521, 65519, 65497],
/// prime_suspects::segmented_sieve(65537, 256)
///   .iter().rev().take(3).map(|&num| num)
///   .collect::<Vec<usize>>());
/// ```
///
// ```
// assert_eq!(&999983,
// prime_suspects::segmented_sieve(1000000, 350).last().unwrap());
// ```
pub fn segmented_sieve(max_val: usize, mut segment_size: usize) -> Vec<usize> {
  if max_val <= ((2 as i64).pow(16) as usize) {
    // early return if the highest value is small enough (empirical)
    return eratosthenes_sieve(max_val);
  }

  if segment_size > ((max_val as f64).sqrt() as usize) {
    segment_size = (max_val as f64).sqrt() as usize;
    println!("Segment size is larger than √{}. Reducing to {} to keep resource use down.",
             max_val, segment_size);
  }

  // get the primes up to the first segment
  let small_primes = eratosthenes_sieve((max_val as f64).sqrt() as usize);
  let mut big_primes = small_primes.clone();

  // As Sorensen says, we need to construct a sequence over each segment, in
  // the interval [start + 1, start + segment_size] that begins with
  // (start + this_prime - ( start mod p)), and increases by p up to
  // (start + segment_size).
  // That sequence will be the values to sieve out of this_segment.

  // clunky way of doing each_slice, from
  // https://stackoverflow.com/a/37033906/2023432
  let mut segment_range = (segment_size..max_val).peekable();
  while segment_range.peek().is_some() {
    let this_segment: Vec<_> = segment_range.by_ref().take(segment_size).collect();
    let mut sieved_segment: Vec<_> = this_segment.clone();
    for &this_prime in &small_primes {
      if !this_segment.is_empty() {
        let mut starting_offset = this_segment[0] % this_prime;
        starting_offset = if starting_offset == 0 { this_prime } else { starting_offset };

        let first_val = this_segment[0] + this_prime - starting_offset;
        let last_val: &usize = this_segment.last().unwrap();

        // hack for inclusive range while RFC is figured out. see
        // https://www.reddit.com/r/rust/comments/3xkfro/what_happened_to_inclusive_ranges/
        let sieve_vec = (first_val..(*last_val + 1))
          .step(this_prime)
          .collect::<Vec<_>>();

        sieved_segment = sieved_segment
          .iter()
          .filter(|&check_num| !sieve_vec.contains(&check_num))
          .map(|&val| val)
          .collect::<Vec<_>>();

      }
    }
    for sieved_prime in sieved_segment {
      big_primes.push(sieved_prime);
    }
  }

  return big_primes;
}

#[test]
fn no_end_segment_sieve_misses() {
  let test_100k_primes = segmented_sieve(100000, 300);
  assert!(!test_100k_primes.contains(&99999));

  let test_100m_primes = segmented_sieve(1000000, 350);
  assert!(test_100m_primes.contains(&999983));
  assert!(!test_100m_primes.contains(&999997));
}
\$\endgroup\$
1
\$\begingroup\$
  1. Documentation of the crate itself should use //! instead of ///.

  2. Documenting "A struct" is redundant and listing where a type is used is overly brittle - will you remember to change that when the other location changes?

  3. Much of the documentation seems redundant in general. Use meaningful variable names, create meaningful types, and don't duplicate the implementation in comments ("starts at ...", "the argument", etc.).

  4. Trailing commas are the prevalent style in Rust.

  5. 4-space indents. It may not be what you are used to, but that's what the community and language creators have chosen. It's better to just accept that and integrate.

  6. Remember that /// is for user-facing documentation, and // is for the programmer. There's no need to tell the user of your code how the Iterator trait works.

  7. It's better to document limitations about the iterator on the type or on the constructor as it will be easier to find that way.

  8. "this specific Iterator will never return None" is not always true. In debug mode, it will abort the program when the value exceeds 64 bits. In a release build, it is true, but the value will wrap back around to zero. This is important to describe to users so they avoid it.

  9. It's great to see you are writing docs and doc-tests, as well as linking to your algorithm sources. Take care to indent your code example blocks the same.

  10. Use Markdown links that place the URL out of the flow of text. This makes the documentation easier to read when viewing the source.

  11. If you have a vector that you no longer need and want to iterate through it by value, use into_iter instead of iter().cloned().

  12. Adding one to the square root seems brittle - what if the value is exact enough that it doesn't round down? Use ceil instead, which is also self-documenting.

  13. Don't compare a boolean value against true or false, just use it directly.

  14. Use Iterator methods like filter and collect instead of loops and mutation. This would create (2..max_val).filter(|&term| bool_vec[term]).collect().

  15. Try to avoid array indexing when possible as it incurs a small penalty for bounds checking. Iterators tend to be more efficient.

  16. It looks like you have a commented-out documentation test. Don't leave commented-out code; fix it or remove it.

  17. I'm unclear as to why you used an i64 to compute the max before converting back to a usize. It seems simpler to just use a usize throughout.

  18. You can use a usize (or i8, f32, etc.) suffix to pick a certain size value instead of as.

  19. Do the floating point / square root / integer conversion once. Besides being slightly more efficient for the computer, it's easier on the programmer to know it's the same value.

  20. Consider rebinding a variable instead of making it mutable. A mutable variable can be accidentally changed easier than being bound again.

  21. There's no need to specify the type on the result of clone.

  22. I'm not a fan of the turbofish operator (::<>) and prefer to specify the type on the variable binding.

  23. map(|&val| val) is called cloned

  24. .filter(|&val| function(&val)) is redundant - why dereference and then re-reference?

  25. There's no need for a type specification when collecting back into sieved_segment, it can only be a Vec.

  26. Use extend instead of a for loop to add more values to a vector (or other methods on Vec that accept bigger chunks).

  27. Don't use explicit returns at the end of functions.

  28. Why pack all your tests into one test method? Split them apart, giving each useful names that distinguish what value each brings.


extern crate itertools;

use itertools::Itertools;


/// Generates the sequence i^2, i^2+i, i^2+2i, i^2+3i, ...
///
/// This `Iterator` does not deal with the sequence value exceeding
/// `usize`; it is the callers' responsibility.
struct SquareMultiple {
    current_value: usize,
    increment: usize,
}

impl Iterator for SquareMultiple {
    type Item = usize;

    fn next(&mut self) -> Option<usize> {
        let val = self.current_value;
        self.current_value += self.increment;
        Some(val)
    }
}

impl SquareMultiple {
    fn new(term: usize) -> Self {
        SquareMultiple { current_value: term * term, increment: term }
    }
}


/// An implementation of the sieve of Eratosthenes, as described in
/// [the Wikipedia article][wiki].
///
/// # Examples
///
/// ```
/// assert_eq!(vec![2,3,5,7], prime_suspects::eratosthenes_sieve(10));
/// ```
///
/// ```
/// assert_eq!(vec![65521, 65519, 65497],
/// prime_suspects::eratosthenes_sieve(65535)
///   .into_iter().rev().take(3)
///   .collect::<Vec<_>>());
/// ```
///
/// [wiki]: https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes#Algorithm_and_variants
pub fn eratosthenes_sieve(max_val: usize) -> Vec<usize> {
    // Algorithm notes: The sieve works like so (Wikipedia, pseudocode inlined):
    // Input: an integer n > 1
    // Let A be an array of Boolean values, indexed by integers 2 to n,
    // initially all set to true.
    let mut bool_vec = vec![true; max_val];
    // for i = 2, 3, 4, ..., not exceeding √n:
    let top_sieve = max_val as f64;
    // We have to add 1 because the sqrt coerced to an int is √floor(n)

    let top_sieve = top_sieve.sqrt().ceil();
    for sieve_term in 2..(top_sieve as usize) {
        // if A[i] is true:
        if bool_vec[sieve_term] {
            // for j = i^2, i^2+i, i^2+2i, i^2+3i, ...,
            for j in SquareMultiple::new(sieve_term)
                .take_while(|&term| term < max_val) { // ...not exceeding n
                    bool_vec[j] = false; // A[j] := false
                }
        }
    }

    // Output: all i such that A[i] is true.
    bool_vec[2..].iter().enumerate().filter(|&(_, &b)| b).map(|(i, _)| i + 2).collect()
}


/// A [segmented approach][wiki] to sieveing, keeping memory use to
/// O(√n). As Sorensen states, this is the most practical optimization
/// to the sieve of Eratosthenes.
///
/// # Examples
///
/// ```
/// assert_eq!(vec![65521, 65519, 65497],
/// prime_suspects::segmented_sieve(65537, 256)
///   .iter().rev().take(3).map(|&num| num)
///   .collect::<Vec<usize>>());
/// ```
///
/// [wiki]: https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes#Segmented_sieve
// ```
// assert_eq!(&999983,
// prime_suspects::segmented_sieve(1000000, 350).last().unwrap());
// ```

pub fn segmented_sieve(max_val: usize, segment_size: usize) -> Vec<usize> {
    if max_val <= 2usize.pow(16) {
        // early return if the highest value is small enough (empirical)
        return eratosthenes_sieve(max_val);
    }

    let alpha = (max_val as f64).sqrt() as usize;

    let segment_size = if segment_size > alpha {
        println!("Segment size is larger than √{}. Reducing to {} to keep resource use down.",
                 max_val, alpha);
        alpha
    } else {
        segment_size
    };

    // get the primes up to the first segment
    let small_primes = eratosthenes_sieve(alpha);
    let mut big_primes = small_primes.clone();

    // As Sorensen says, we need to construct a sequence over each
    // segment, in the interval [start + 1, start + segment_size] that
    // begins with (start + this_prime - ( start mod p)), and
    // increases by p up to (start + segment_size).
    //
    // That sequence will be the values to sieve out of this_segment.

    // clunky way of doing each_slice, from
    // http://stackoverflow.com/a/37033906/2023432
    let mut segment_range = (segment_size..max_val).peekable();
    while segment_range.peek().is_some() {
        let this_segment: Vec<_> = segment_range.by_ref().take(segment_size).collect();


        let mut sieved_segment = this_segment.clone();
        for &this_prime in &small_primes {
            if !this_segment.is_empty() {
                let mut starting_offset = this_segment[0] % this_prime;
                starting_offset = if starting_offset == 0 { this_prime } else { starting_offset };

                let first_val = this_segment[0] + this_prime - starting_offset;
                let last_val: &usize = this_segment.last().unwrap();

                // hack for inclusive range while RFC is figured out. see
                // https://www.reddit.com/r/rust/comments/3xkfro/what_happened_to_inclusive_ranges/

                let sieve_vec: Vec<_> = (first_val..(*last_val + 1))
                    .step(this_prime)
                    .collect();

                sieved_segment = sieved_segment
                    .iter()
                    .filter(|check_num| !sieve_vec.contains(check_num))
                    .cloned()
                    .collect();

            }
        }

        big_primes.extend(sieved_segment);
    }

    big_primes
}

#[test]
fn no_end_segment_sieve_misses() {
  let test_100k_primes = segmented_sieve(100000, 300);
  assert!(!test_100k_primes.contains(&99999));

  let test_100m_primes = segmented_sieve(1000000, 350);
  assert!(test_100m_primes.contains(&999983));
  assert!(!test_100m_primes.contains(&999997));
}
\$\endgroup\$
  • 1
    \$\begingroup\$ I hope you get a visit from Veedrac - they are amazing at analyzing algorithms and making them better; I mostly did stylistic and idiomatic concerns. I'm pretty sure that the whole each_slice can be simplified. \$\endgroup\$ – Shepmaster May 18 '16 at 0:52
  • \$\begingroup\$ Thanks; the idiom/style issues were the biggest concern for me. \$\endgroup\$ – bright-star May 18 '16 at 1:18

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.