Implementing a custom iterator

To understand iterators more thoroughly, we'll implement an iterator that generates prime numbers up to a certain limit that's customizable by the user. First, let's clarify the API expectations that we'll need from our iterator:

// custom_iterator.rs

use std::usize;

struct Primes {
limit: usize
}

fn main() {
let primes = Primes::new(100);
for i in primes.iter() {
println!("{}", i);
}
}

So, we have a type called Primes that we can instantiate with the new method, providing the upper bound on the number of primes to generate. We can call iter() on this instance to convert it in to an iterator type, which can then be used in a for loop. With that said, let's add the new and iter methods on it:

// custom_iterator.rs

impl Primes {
fn iter(&self) -> PrimesIter {
PrimesIter {
index: 2,
computed: compute_primes(self.limit)
}
}

fn new(limit: usize) -> Primes {
Primes { limit }
}
}

The iter method takes the Primes type via &self and returns a PrimesIter type containing two fields: index, which stores the index in the vector, and a computed field that stores the pre-computed primes in a vector. The compute_primes method is defined as follows:

// custom_iterator.rs

fn compute_primes(limit: usize) -> Vec<bool> {
let mut sieve = vec![true; limit];
let mut m = 2;
while m * m < limit {
if sieve[m] {
for i in (m * 2..limit).step_by(m) {
sieve[i] = false;
}
}
m += 1;
}
sieve
}

This function implements the sieve of the eratosthenes algorithm for efficiently generating prime numbers up to a given limit. Next, there's the definition of the PrimesIter struct along with its Iterator implementation:

// custom_iterator.rs

struct PrimesIter {
index: usize,
computed: Vec<bool>
}

impl Iterator for PrimesIter {
type Item = usize;
fn next(&mut self) -> Option<Self::Item> {
loop {
self.index += 1;
if self.index > self.computed.len() - 1 {
return None;
} else if self.computed[self.index] {
return Some(self.index);
} else {
continue
}
}
}
}

In the next method, we loop and get the next prime number if the value at self.index is true in the self.computed Vec. If we went past the elements in our computed Vec, then we return None to signify that we are done. Here's the complete code with the main function that generates 100 prime numbers:

// custom_iterator.rs

use std::usize;

struct Primes {
limit: usize
}

fn compute_primes(limit: usize) -> Vec<bool> {
let mut sieve = vec![true; limit];
let mut m = 2;
while m * m < limit {
if sieve[m] {
for i in (m * 2..limit).step_by(m) {
sieve[i] = false;
}
}
m += 1;
}
sieve
}

impl Primes {
fn iter(&self) -> PrimesIter {
PrimesIter {
index: 2,
computed: compute_primes(self.limit)
}
}

fn new(limit: usize) -> Primes {
Primes { limit }
}
}

struct PrimesIter {
index: usize,
computed: Vec<bool>
}

impl Iterator for PrimesIter {
type Item = usize;
fn next(&mut self) -> Option<Self::Item> {
loop {
self.index += 1;
if self.index > self.computed.len() - 1 {
return None;
} else if self.computed[self.index] {
return Some(self.index);
} else {
continue
}
}
}
}

fn main() {
let primes = Primes::new(100);
for i in primes.iter() {
print!("{},", i);
}
}

We get the following output:

3,5,7,11,13,17,19,23,29,31,37,41,43,47,53,59,61,67,71,73,79,83,89,97

Great! Apart from Vec, there are a lot of types that implement the Iterator trait in the standard library, such as HashMap, BTreeMap, and VecDeque.

..................Content has been hidden....................

You can't read the all page of ebook, please click here login for view all page.
Reset