22-实现自定义迭代器

实现自定义迭代器

为了更深入地了解迭代器，我们将实现一个自定义迭代器，它能够生成一定数量的素数。首先，让我们来了解一下迭代器中需要用到的API：

// custom_iterator.rs
use std::usize;
struct Primes {
    limit: usize
}
fn main() {
    let primes = Primes::new(100);
    for i in primes.iter() {
        println!("{}", i);
    }
}

我们有一个名为Primes的类型，可以使用new方法并提供生成素数的数目来实例化它。我们可以在这个实例上调用iter()将其转换为迭代器类型，然后在for循环中使用它。接下来，让我们给它添加new和iter方法：

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

iter方法通过&self获取Primes类型，并返回包含两个字段的PrimesIter类型：index用于将索引值存储在 vector 中，computed 字段用于将预先计算的素数存储在 vector 中。compute_primes方法的定义如下所示：

// 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
}

该函数实现了素数筛选算法（sieve of the eratosthenes algorithm），用于高效地生成给定数目的素数。接下来是PrimesIter结构体定义，以及Iterator的实现：

// 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
            }
        }
    }
}

在next方法中，我们执行循环，如果self.index对应的self.computed中的值为true，那么获取其中的素数。如果我们遍历了computed容器中的元素，那么返回None，表示迭代访问操作已完成。这是包含main函数的完整代码，它会生成100个素数：

// 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);
    }
}

我们得到以下输出结果：

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

除了Vec之外，还有很多类型在标准库中实现了Iterator特征，例如HashMap、BtreeMap及VecDeque。