Rust Iterator

1. Rust 迭代器
- 1.1. 3 种形式的迭代器（iter/iter_mut/into_iter）
- 1.2. for 循环是迭代器的语法糖
2. 实现迭代器
3. 迭代器是惰性的
4. 迭代器适配器（一个迭代器转换为另一个迭代器）
5. 迭代器消费者（触发遍历）
6. 第三方库
- 6.1. itertools
- 6.2. rayon
7. 参考

1. Rust 迭代器

Rust 中使用迭代器可以方便地处理序列数据。

迭代器需要实现 Iterator trait：

trait Iterator {
    type Item;
    fn next(&mut self) -> Option<Self::Item>;
}

从上定义可知，迭代器只有一个必须实现的方法，即 next() 。

Rust 中的 Range 就实现了 Iterator trait，下面是它的例子：

fn main() {
    let mut x = 0..3;

    assert_eq!(x.next(), Some(0));
    assert_eq!(x.next(), Some(1));
    assert_eq!(x.next(), Some(2));
    assert_eq!(x.next(), None);
}

又如， String 和 &str 的 lines() 方法返回的也是迭代器，下面是它的例子：

fn main() {
    let text = "line 1\nline 2";
    let mut x = text.lines();           // lines() 返回一个迭代器

    assert_eq!(x.next(), Some("line 1"));
    assert_eq!(x.next(), Some("line 2"));
    assert_eq!(x.next(), None);
}

1.1. 3 种形式的迭代器（iter/iter_mut/into_iter）

一般地，从一个 Collection 上的可以得到 3 种形式的迭代器（有些 Collection 可能只实现了其中 1 种或者 2 种）：

iter(), which iterates over &T.
iter_mut(), which iterates over &mut T.
into_iter(), which iterates over T.

下面依次介绍它们。

iter() 最为常用，下面是它的例子：

fn main() {
    let x = [1, 2, 3];
    for element in x.iter() {           // 使用 iter() 得到迭代器
        println!("{}", element);        // element 的类型是 &i32
    }
}

通过 iter_mut() 可以修改 Collection 中数据，下面是它的例子：

fn main() {
    let mut x = [1, 2, 3];

    for element in x.iter_mut() {       // 使用 iter_mut() 得到可修改迭代器
        *element += 1;                  // 每个元素都加 1，element 的类型是 &mut i32
    }

    println!("{:?}", x);                // [2, 3, 4]
}

使用 into_iter() 得到的迭代器在遍历时会消耗掉 Collection 中元素，下面是它的例子：

fn main() {
    #[derive(Debug)]
    struct Number(i32);

    let x = vec![Number(1), Number(2), Number(3)];
    for element in x.into_iter() {    // Number 未实现 Copy trait，所有权会转移到 element 上
        println!("{:?}", element)     // element 的类型是 Number，如果上一行使用 x.iter()，则 element 的类型是 &Number
    }

    // println!("{:?}", x);           // 这行会报错。x 中所有元素都已经被消耗了，不再可访问。如果前面改为 x.iter() 则还可访问
}

上面场景中， into_iter() 可以直接省略，也就是说上面代码等同于：

fn main() {
    #[derive(Debug)]
    struct Number(i32);

    let x = vec![Number(1), Number(2), Number(3)];
    for element in x {                // for ... in x 等同于 for ... in x.into_iter()
        println!("{:?}", element)
    }

    // println!("{:?}", x);
}

1.2. for 循环是迭代器的语法糖

Rust 中 for 循环是迭代器的语法糖，比如下面 for 循环：

let values = vec![1, 2, 3, 4, 5];

for x in values {
    println!("{}", x);
}

相当于：

let values = vec![1, 2, 3, 4, 5];
{
    let result = match IntoIterator::into_iter(values) {
        mut iter => loop {
            let next;
            match iter.next() {           // 调用迭代器的 next() 方法
                Some(val) => next = val,
                None => break,
            };
            let x = next;
            let () = { println!("{}", x); };
        },
    };
    result
}

2. 实现迭代器

下面是为 Counter 实现迭代器的例子，从 1 到 5 计数：

// First, the struct:
struct Counter {
    count: usize,
}

// we want our count to start at one, so let's add a new() method to help.
// This isn't strictly necessary, but is convenient. Note that we start
// `count` at zero, we'll see why in `next()`'s implementation below.
impl Counter {
    fn new() -> Counter {
        Counter { count: 0 }
    }
}

// Then, we implement `Iterator` for our `Counter`:
impl Iterator for Counter {
    // we will be counting with usize
    type Item = usize;

    // next() is the only required method
    fn next(&mut self) -> Option<Self::Item> {
        // Increment our count. This is why we started at zero.
        self.count += 1;

        // Check to see if we've finished counting or not.
        if self.count < 6 {
            Some(self.count)
        } else {
            None
        }
    }
}

fn main() {
    let mut counter = Counter::new();

    assert_eq!(counter.next(), Some(1));
    assert_eq!(counter.next(), Some(2));
    assert_eq!(counter.next(), Some(3));
    assert_eq!(counter.next(), Some(4));
    assert_eq!(counter.next(), Some(5));
    assert_eq!(counter.next(), None);
}

3. 迭代器是惰性的

迭代器是惰性的，创建迭代器时并不会对 Collection 中元素进行遍历，直到遇到迭代器消费者（它会调用迭代器的 next() 方法）才会触发遍历。 比如：

fn main() {
    let text = " line 1  \n line 2  ";
    let trimmed = text.lines()              // lines() 返回一个迭代器，现在还不会开始遍历 text
        .map(str::trim)                     // map 返回另一个迭代器，现在还不会开始遍历 text
        .collect::<Vec<&str>>();            // collect() 是迭代器消费者，这时才开始遍历 text

    println!("{:?}" , trimmed);             // ["line 1", "line 2"]
}

4. 迭代器适配器（一个迭代器转换为另一个迭代器）

通过迭代器适配器，可以把一个迭代器转换为另一个迭代器，比如 map, filter 是常用的迭代器适配器。

4.1. map

在节 3 中介绍过 map，它可以对 Collection 中元素按规则进行映射，下面再看个 map 的使用例子：

fn main() {                                 // iter().map 形式一
    let a: Vec<i32> = vec![1, -2, -3, 4, 5];
    let b = a
        .iter()
        .map(|x: &i32| (*x) * 3);           // x 类型为 &i32，这是因为 iter() 在类型 &T 上遍历，(*x) * 3 可简化，后面会介绍
    let c: Vec<i32> = b.collect();

    println!("{:?}", c)                     // [3, -6, -9, 12, 15]
}

上面例子中，map 接收 closure，这个 closure 的参数类型是 &i32 ，为什么是这样呢？这是因为 iter() 会在 &T 上遍历，详情可参考节 1.1 。

如果我们把上面例子中的 iter() 换为 into_iter() ，则 closure 的参数类型就是 i32 了：

fn main() {
    let a: Vec<i32> = vec![1, -2, -3, 4, 5];
    let b = a
        .into_iter()
        .map(|x: i32| x * 3);               // x 类型为 i32，这是因为 into_iter() 在类型 T 上遍历
    let c: Vec<i32> = b.collect();

    println!("{:?}", c)                     // [3, -6, -9, 12, 15]
}

需要说明的是，由于类型 &i32 “重载”了乘法操作符 Mul ，可以直接把 &i32 和 i32 相乘：

fn main() {                                 // iter().map 形式二
    let a: Vec<i32> = vec![1, -2, -3, 4, 5];
    let b = a
        .iter()
        .map(|x: &i32| x * 3);              // x 类型为 &i32，也可以直接把它和数字 3 相乘
    let c: Vec<i32> = b.collect();

    println!("{:?}", c)                     // [3, -6, -9, 12, 15]
}

fn main() {                                 // iter().map 形式三
    let a: Vec<i32> = vec![1, -2, -3, 4, 5];
    let b = a
        .iter()
        .map(|&x| x * 3);                   // closure 参数类型为 &i32，如果参数写为 &x，则 x 类型就是 i32
    let c: Vec<i32> = b.collect();

    println!("{:?}", c)                     // [3, -6, -9, 12, 15]
}

4.2. filter

filter 可以对 Collection 中元素按规则进行过滤，下面是 filter 的使用例子：

fn main() {
    let a: Vec<i32> = vec![1, -2, -3, 4, 5];
    let b = a
        .iter()                                 // iter() 会在 &T 上遍历
        .filter(|x: &&i32| **x > 0);            // x 的类型为 &&i32，一个 & 源于 iter()，另一个 & 源于 filter
    let c: Vec<&i32> = b.collect();

    println!("{:?}", c)                         // [1, 4, 5]
}

这里需要注意一下，上面例子中 filter 参数是 closure，closure 的参数 x 的类型为 &&i32 ，这导致 closure 内部使用 **x 才能得到原始的值。

我们知道，使用 iter() 时，会在 &T 上遍历，之所以还多一个 & 的原因是 filter 本身就接收一个引用，下面是 filter 的签名：

fn filter<P>(self, predicate: P) -> Filter<Self, P>
where
    Self: Sized,
    P: FnMut(&Self::Item) -> bool
//           ^
//           |
//       这里有个 &

如果我们把上面例子中的 iter() 换为 into_iter() ，则 closure 的参数类型就是 &i32 了：

fn main() {
    let a: Vec<i32> = vec![1, -2, -3, 4, 5];
    let b = a
        .into_iter()                            // into_iter() 会在 T 上遍历
        .filter(|x: &i32| *x > 0);              // x 的类型为 &i32，它源于 filter
    let c: Vec<i32> = b.collect();

    println!("{:?}", c)                         // [1, 4, 5]
}

为了让 closure 内部代码简单些，当使用 iter() 时，可以使用下面不同的形式：

    let a: Vec<i32> = vec![1, -2, -3, 4, 5];
    let mut iter = a.iter().filter(|x| **x > 1);      // 形式一
    let mut iter = a.iter().filter(|&x| *x > 1);      // 同上，形式二
    let mut iter = a.iter().filter(|&&x| x > 1);      // 同上，形式三

当使用 into_iter() 时，可以使用下面不同的形式：

    let a: Vec<i32> = vec![1, -2, -3, 4, 5];
    let mut iter = a.into_iter().filter(|x| *x > 1);  // 形式一
    let mut iter = a.into_iter().filter(|&x| x > 1);  // 同上，形式二

4.3. filter_map

filter_map 可以看作是 filter 和 map 的混合体，可以使代码更加简洁。

比如， 我们有一个字符串数组，现在要过滤出那些可以转换为数字的，并转换为数字。 这个任务可以通过先 map，再 filter，再 map 来实现：

fn main() {
    let a = ["1", "two", "NaN", "four", "5"];
    let mut iter = a.iter()
        .map(|s| s.parse())                   // 字符串转换为数字，提到 Result
        .filter(|s| s.is_ok())                // 过滤可以 parse 成功的
        .map(|s| s.unwrap());                 // 上一步已经过滤了，直接 unwrap 可得到数字，不会失败

    assert_eq!(iter.next(), Some(1));
    assert_eq!(iter.next(), Some(5));
    assert_eq!(iter.next(), None);
}

如果使用 filter_map ，则上面代码可以简化为：

fn main() {
    let a = ["1", "two", "NaN", "four", "5"];
    let mut iter = a.iter()
        .filter_map(|s| s.parse().ok());

    assert_eq!(iter.next(), Some(1));
    assert_eq!(iter.next(), Some(5));
    assert_eq!(iter.next(), None);
}

ok() 表示如果 Result 是成功的，返回 Some(X)，否则返回 None；而 filter_map 会从 closure 返回结果中过滤出 Some(X)，返回 X。

4.4. flatten

flatten 可以拍平嵌套的迭代器，下面是 flatten 的一个使用实例：

fn main() {
    let data = vec![vec![1, 2, 3, 4], vec![5, 6]];
    let flattened = data.into_iter().flatten().collect::<Vec<u8>>();

    assert_eq!(flattened, &[1, 2, 3, 4, 5, 6]);
}

4.5. flat_map

flat_map 的参数是个 closure，这个 closure 返回一个迭代器，而 flat_map 会把 closure 返回的多个迭代器拍平 ，下面是个例子：

fn main() {
    let lines_vec = vec!("hello:rust","how:are:you");
    let words_vec = lines_vec.iter()
        .flat_map(|x: &&str| (*x).split(":"))   // split 返回的是迭代器
        .collect::<Vec<&str>>();

    println!("{:?}", words_vec);                // ["hello", "rust", "how", "are", "you"]
}

下面是 flat_map 的另一个使用例子：

use std::collections::HashMap ;

fn main() {
    let mut scores = HashMap::new() ;

    scores.insert("alice", vec![70, 50, 90]);
    scores.insert("bob", vec![70, 80, 90]);
    scores.insert("carol", vec![40, 60, 70]);

    let people = ["alice", "bob", "carol"];

    let x: Vec<&i32> = people.iter()
        .flat_map(|x| scores[x].iter())         // closure 返回迭代器，flat_map 把所有迭代器拍平
        .collect();

    println!("{:?}", x);                        // [70, 50, 90, 70, 80, 90, 40, 60, 70]
}

再看一个 flat_map 的使用例子：

fn main() {
    let mut scores = HashMap::new() ;

    scores.insert("alice", vec![70, 50, 90]);
    scores.insert("bob", vec![70, 80, 90]);
    scores.insert("carol", vec![40, 60, 70]);

    let people = ["alice", "bob", "carol"];

    let x: Vec<i32> = people.iter()
        .flat_map(|x| scores[x].iter().map(|y| y + 1))  // closure 返回迭代器，flat_map 把所有迭代器拍平
        .collect();

    println!("{:?}", x);                        // [71, 51, 91, 71, 81, 91, 41, 61, 71]
}

4.6. 其它迭代器适配器

还有很多这里未介绍的迭代器适配器，可参考：

5. 迭代器消费者（触发遍历）

创建迭代器时不会遍历 Collection，只有遇到迭代器消费者时，才会触发遍历。常见的迭代器消费者有 collect, fold, count 等等。

5.1. collect

在 4 中已经介绍了 collect 的很多使用例子。使用 collect 可以把遍历迭代器时的数量收集到一个容器中。不过，我们需要指定你想收集到什么容器中，比如 Vec 或者 LinkedList 等，下面是使用例子：

use std::collections::LinkedList;

fn main() {
    let a: [i32; 3] = [1, 2, 3];

    let doubled: Vec<i32> = a.iter().map(|&x| x * 2).collect();
    let doubled: LinkedList<i32> = a.iter().map(|&x| x * 2).collect();

    let doubled = a.iter().map(|&x| x * 2).collect::<Vec<i32>>();
    let doubled = a.iter().map(|&x| x * 2).collect::<LinkedList<i32>>();

    // 省略元素类型
    let doubled: Vec<_> = a.iter().map(|&x| x * 2).collect();
    let doubled: LinkedList<_> = a.iter().map(|&x| x * 2).collect();

    let doubled = a.iter().map(|&x| x * 2).collect::<Vec<_>>();
    let doubled = a.iter().map(|&x| x * 2).collect::<LinkedList<_>>();
}

在指定类型时，编译器仅想知道哪种类型的容器（比如是 Vec 还是 LinkedList 或者其它），而容器元素类型（上面例子中的 i32）是可以推导出来的，所以不需要显示指定元素类型，如上面例子中的 Vec<i32> 可以简写为 Vec<_> 。

5.2. fold（Reduce 操作）

fold 用于从 Collection 数据得到单个数据，这也称为 Reduce 操作。 它接收 2 个参数，第 1 个参数是初始值，第 2 个参数是一个 closure。

下面是使用 fold 求和的例子（真实项目中直接使用 sum 即可，这里仅仅是演示 fold 用法）：

fn main() {
    let a = [1, 2, 3, 4];

    // the sum of all of the elements of the array
    let sum = a.iter().fold(0, |acc: i32, x: &i32| acc + x);
    //                      ^     ^       ^
    //                      |     |       |
    //                   initial  |    element
    //                            |
    //                        accumulator

    assert_eq!(sum, 10);
}

表 1 演示了上面代码执行时，closure 两个参数 acc 和 x 的变化情况。

Table 1: a.iter().fold(0, |acc, x| acc + x) 执行时 acc 和 x 的变化情况
element	acc	x	result
	0
1	0	1	1
2	1	2	3
3	3	3	6
4	6	4	10

下面是使用 fold 分别实现 sum, count, product 的例子：

fn main() {
    let a = [1, 2, 3, 4];

    println!("sum {}", a.iter().fold(0, |acc, x| acc + x));       // sum 10
    println!("count {}", a.iter().fold(0, |acc, _| acc + 1));     // count 4
    println!("product {}", a.iter().fold(1, |acc, x| acc * x));   // product 24
}

5.3. 其它迭代器消费者

还有很多这里未介绍的迭代器消费者，可参考：

6. 第三方库

6.1. itertools

标准库已经提供了很多迭代器适配器和迭代器消费者，如果还不够用，可以看看 itertools。

6.2. rayon

如果你想要个并行库，可以使用 rayon 。把标准库中节 1.1 介绍的迭代器换为 rayon 中如表 2 所示的对应迭代器即可实现并行操作。

Table 2: rayon 并发库 3 种形式的迭代器
标准库	rayon 并发库
iter()	par_iter()
iter_mut()	par_iter_mut()
into_iter()	into_par_iter()

下面是 rayon 的一个例子：

use rayon::prelude::*;

fn sum_of_squares(input: &[i32]) -> i32 {
    input.par_iter()                // <-- just change that!
         .map(|&i| i * i)
         .sum()
}