MashPlant的笔记

“你将孤单度过一生”

0%

试用Rust的Generator

Posted on Edited on Views: Valine:

简单记录一下我学习和使用Rust的Generator的过程。

背景

Rust中的Generator这一套很久之前就提出来了，在issue#43122里一直从2017年讨论到现在，还是没有稳定，我也没见到什么知名的项目用过它，如果是我见识太短浅了麻烦提醒我。好像这个东西就一直游离在语言的边缘，看上去有点用，但是一直没有人用。我不太了解历史，但是总觉得async那一套出现的比它还晚，成熟的比它还早，如果我错了的话也麻烦提醒一下。

Generator trait

类比async的核心是Future trait，Generator的核心是Generator trait。把这两个trait的定义列出来比较一下：

1
2
3
4
5
6
7
8
9
10
pub trait Future {
  type Output;
  fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output>;
}

pub trait Generator<R = ()> {
  type Yield;
  type Return;
  fn resume(self: Pin<&mut Self>, arg: R) -> GeneratorState<Self::Yield, Self::Return>;
}

其实是相当类似的。简单描述一下相同点和不同点：

  1. 都以self: Pin<&mut Self>为参数，因为它们内部都可能保存了指向自身的指针，只有这个数据结构被Pin住的情况下操作它才是安全的。
  2. Future只有一个类型参数Output，表示最终返回的结果，Poll其实就相当于Option；Generator有两个类型参数Yield和Return，可以在返回零次或多次Yield后最终返回Return，GeneratorState其实就相当于Result。
  3. Future的poll中有一个额外参数cx，一般是用来注册回调的；Generator的resume中有一个额外参数arg，也就是每次进入的时候都可以传一个参数，但是我个人觉得这和我印象中的生成器并不是很匹配，生成器不应该是给出一个初始参数构造好之后就像一个迭代器一样不断生成值吗？每次生成值的时候传入新的参数，这是一个什么样的语义呢？什么情况下能用得到呢？我现在也还不理解，希望得到解答。

用例

unstable-book中有一章讲Generator，我稍微简化了一下它的例子：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
#![feature(generators, generator_trait)]

use std::ops::{Generator, GeneratorState};
use std::pin::Pin;

fn main() {
  let mut g = || {
    yield 1;
    return "foo"
  };

  assert_eq!(Pin::new(&mut g).resume(()), GeneratorState::Yielded(1));
  assert_eq!(Pin::new(&mut g).resume(()), GeneratorState::Complete("foo"));
}

就像一般不会手动实现Future trait，而是用async fn或者async块来得到实现了Future trait的匿名类型一样，一般也不会手动实现Generator trait。上面的代码中g的值是一个generator literal，形式是一个闭包，里面可以用yield关键字来yield一个值。g的类型就是一个实现了Generator<(), Yield=i32, Return=&'static str>的匿名类型。

为了调用g，需要先用Pin把它包裹起来再调用resume，这是Generator trait的要求。每次调用resume，就从上次停下来的地方开始(第一次执行时就从开头开始)，执行到yield或者return或者末尾，从而yield或者return一个值。这还是很符合直觉的，其他语言中如果有类似的feature的话，语义差不多也是这样的。

需要注意的是，这里用了闭包的形式来定义Generator，这并不是为了简化代码而不用函数定义，事实上当前的编译器中不支持用函数定义Generator。具体细节请看error#E0627。而Rust中的闭包想要引用自身相当麻烦（确实是可行的，但是需要一堆丑陋的boilerplate），所以目前应该是不太可能像其他语言（比如Python）一样很愉快地写递归的Generator。

意义

unstable-book中说：

At this time the main intended use case of generators is an implementation primitive for async/await syntax, but generators will likely be extended to ergonomic implementations of iterators and other primitives in the future.

这里面说了两点，第一是作为实现async那一套的原语，我理解可能是编译器先把async那一套desugar成yield那一套，这样后续步骤就只用处理yield，毕竟这两种东西的语义还是有很多近似的地方，这样做是可以理解的，不过跟我们用户好像没什么关系。

第二是用来实现迭代器之类的，这个我觉得是一个很符合实际需求的用途。比如我可能有一个递归函数，里面产生一些值，希望能把值的序列以迭代器的形式返回给用户，在当前的稳定Rust中我能想到的唯一的方法就是自己手动维护栈，实现状态机。如果能够在产生值的地方yield，然后整个函数自动成为一个Generator，那就方便多了。Generator和Iterator也很类似，前者一直返回Yield直到Return，后者一直返回Some(Item)直到None，所以用户很容易就可以把这个Generator当成迭代器用。

但是现在的Generator是不能直接这样写的，上面也说了，当前的Rust没有提供直接在Generator中递归的手段。

递归

上面说了不能直接写，但是经过一些不算太复杂的改写，还是可以用同样结构的代码来实现这个需求的。

首先造个简单的轮子，把Generator<(), Yield>当Iterator<Item=Yield>用，这里Generator的resume的参数（默认）是单元类型，和我上面说的一样，这才是比较符合直觉的，后续执行的时候应该不用再提供信息；resume的返回值直接丢弃，对应于next返回None。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
#![feature(generators, generator_trait)]

use std::ops::{Generator, GeneratorState};
use std::pin::Pin;

struct Iter<G>(G);

impl<G: Generator + Unpin> Iterator for Iter<G> {
  type Item = G::Yield;

  fn next(&mut self) -> Option<Self::Item> {
    match Pin::new(&mut self.0).resume(()) {
      GeneratorState::Yielded(x) => Some(x),
      GeneratorState::Complete(_) => None,
    }
  }
}

树的中序遍历

先看一个Python版本的：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
class Tree:
  def __init__(self, x, l=None, r=None):
    self.x = x
    self.l = l
    self.r = r

  def inorder(self):
    if self.l: yield from self.l.inorder()
    yield self.x
    if self.r: yield from self.r.inorder()

tree = Tree(
  5,
  Tree(3, Tree(1, None, Tree(2)), Tree(4)),
  Tree(7, Tree(6), Tree(8, None, Tree(9))),
)

print(list(tree.inorder())) # 输出[1, 2, 3, 4, 5, 6, 7, 8, 9]

Python大家应该都会，没有什么可说的。唯一可能不太常用的是yield from，请看pep-380。

先在Rust里造一棵一样的树：

1
2
3
4
5
6
7
8
9
10
11
struct Tree(u64, Option<Box<Tree>>, Option<Box<Tree>>);

fn main() {
  fn inner(x: u64, l: Option<Box<Tree>>, r: Option<Box<Tree>>) -> Option<Box<Tree>> { Some(Box::new(Tree(x, l, r))) }
  fn leaf(x: u64) -> Option<Box<Tree>> { Some(Box::new(Tree(x, None, None))) }
  let tree = Tree(
    5,
    inner(3, inner(1, None, leaf(2)), leaf(4)),
    inner(7, leaf(6), inner(8, None, leaf(9))),
  );
}

我希望不用写闭包引用自身的那一坨代码（其实我也没有试过这个方法是否能用于Generator），所以还是需要一个函数，这个函数自身不可能是Generator，所以只能让它返回一个Generator。先写个大概的框架（相信大家能够理解'a在这里的意义）：

1
2
3
4
5
6
7
8
9
impl Tree {
  fn inorder<'a>(&'a self) -> impl Generator<Yield=u64> + 'a {
    move || {
      if let Some(l) = &self.1 { ??? }
      yield self.0;
      if let Some(r) = &self.2 { ??? }
    }
  }
}

之前说希望能返回给用户一个迭代器用，而不是生成器，所以再改改，用上之前造的Iter：

1
2
3
4
5
6
7
fn inorder<'a>(&'a self) -> impl Iterator<Item=u64> + 'a {
  Iter(move || {
    if let Some(l) = &self.1 { ??? }
    yield self.0;
    if let Some(r) = &self.2 { ??? }
  })
}

问号的地方，就是希望执行Python中的yield from的操作，也就是从一个生成器（在我们这里是迭代器）中生成所有元素，Rust里没有yield from这个集成的操作，只能自己手动写这个逻辑：

1
2
3
4
5
6
7
8
9
10
11
12
13
fn inorder<'a>(&'a self) -> impl Iterator<Item=u64> + 'a {
  Iter(move || {
    if let Some(l) = &self.1 {
      let l = l.inorder();
      for x in l { yield x; }
    }
    yield self.0;
    if let Some(r) = &self.2 {
      let r = r.inorder();
      for x in r { yield x; }
    }
  })
}

这时编译会报错，大概长这个样子：

1
2
3
error[E0720]: cannot resolve opaque type
fn inorder<'a>(&'a self) -> impl Iterator<Item=u64> + 'a {
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ recursive opaque type

虽然不知道具体原因，但是可以猜测是因为递归调用的返回类型与函数的返回类型一样，编译器类型推断的时候出现了循环。这种情况的解决方案一般是用Box<dyn trait>套一层来打破循环：

1
2
3
4
5
6
7
8
9
10
11
12
13
fn inorder<'a>(&'a self) -> impl Iterator<Item=u64> + 'a {
  Iter(move || {
    if let Some(l) = &self.1 {
      let l: Box<dyn Iterator<Item=u64>> = Box::new(l.inorder());
      for x in l { yield x; }
    }
    yield self.0;
    if let Some(r) = &self.2 {
      let r: Box<dyn Iterator<Item=u64>> = Box::new(r.inorder());
      for x in r { yield x; }
    }
  })
}

到这里就已经完成了，加点输出：

1
2
3
4
fn main() {
  ...
  for x in tree.inorder() { print!("{} ", x); } // 输出1 2 3 4 5 6 7 8 9 
}

可以再定义一个宏来提取一下这个手动的yield from的操作：

1
2
3
4
5
6
macro_rules! yield_from {
  ($iter: expr, $elem: ty) => {
    let iter: Box<dyn Iterator<Item=$elem>> = Box::new($iter);
    for x in iter { yield x; }
  };
}

展开列表

这个例子和上面那个没有什么本质区别，可以跳过。

我记得我在初学Python的时候就接触过一个经典的例子：展开一个不规则的列表的列表，比如把[[[1, 2, 3], [4, 5]], 6]展开成[1, 2, 3, 4, 5, 6]：

1
2
3
4
5
6
def flatten(l):
  for x in l:
    if isinstance(x, list): yield from flatten(x)
    else: yield x

print(list(flatten([[[1, 2, 3], [4, 5]], 6]))) # 输出[1, 2, 3, 4, 5, 6]

用Rust来写一下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
enum Elem {
  List(Vec<Elem>),
  Single(u64),
}

fn flatten<'a>(l: &'a [Elem]) -> impl Iterator<Item=u64> + 'a {
  Iter(move || {
    for x in l {
      match x {
        Elem::List(x) => yield_from!(flatten(x), u64),
        Elem::Single(x) => yield *x,
      }
    }
  })
}

fn main() {
  let l = vec![
    Elem::List(vec![
      Elem::List(vec![Elem::Single(1), Elem::Single(2), Elem::Single(3)]),
      Elem::List(vec![Elem::Single(4), Elem::Single(5)]),
    ]),
    Elem::Single(6),
  ];
  for x in flatten(&l) { print!("{} ", x); } // 输出1 2 3 4 5 6 
}

爆栈

需要声明的是，上面这些转化并没有把一个依赖调用栈的递归函数的变成一个依赖堆上的栈进行状态转移的非递归的数，所以该爆栈的时候还是会爆栈的。这在Python里也是一样的，下面这两个例子都会爆栈，如果你的环境下没有爆的话，要么是数字还不够大，要么是编译器做了优化，本质逻辑上还是依赖于一个调用栈的。

1
2
3
4
5
6
def range(n):
  if n > 0: yield from range(n - 1)
  yield n

print(list(range(10))) # 输出[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
print(list(range(1000000))) # 爆栈
1
2
3
4
5
6
7
8
9
10
11
fn range(n: u64) -> impl Iterator<Item=u64> {
  Iter(move || {
    if n > 0 { yield_from!(range(n - 1), u64); }
    yield n;
  })
}

fn main() {
  for x in range(10) { print!("{} ", x); } // 输出0 1 2 3 4 5 6 7 8 9 10 
  for x in range(1000000) { print!("{} ", x); } // 爆栈
}

性能

纯粹无聊，想测试一下几种方法来实现递归的性能怎么样。只是随便测一下，满足我自己的好奇心，并不保证什么严谨性。这里用树的中序遍历来测试，求所有节点的和，实现递归的方法包括直接递归，用Generator实现递归，用标准的栈实现中序遍历的算法实现递归，用栈+状态机实现递归。代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
impl Tree {
  fn inorder_rec(&self, f: &mut impl FnMut(u64)) {
    if let Some(l) = &self.1 { l.inorder_rec(f); }
    f(self.0);
    if let Some(r) = &self.2 { r.inorder_rec(f); }
  }

  fn inorder_gen<'a>(&'a self) -> impl Iterator<Item=u64> + 'a {
    Iter(move || {
      if let Some(l) = &self.1 { yield_from!(l.inorder_gen(), u64); }
      yield self.0;
      if let Some(r) = &self.2 { yield_from!(r.inorder_gen(), u64); }
    })
  }

  // 标准的栈实现中序遍历的算法，上过数据结构课的人应该都能理解
  fn inorder_stk(&self, f: &mut impl FnMut(u64)) {
    let mut stk = vec![self];
    while let Some(&t) = stk.last() {
      let mut t = t;
      while let Some(l) = &t.1 {
        stk.push(l);
        t = l;
      }
      while let Some(t) = stk.pop() {
        f(t.0);
        if let Some(r) = &t.2 {
          stk.push(r);
          break;
        }
      }
    }
  }

  // 栈+状态机实现递归。Rust中没有goto，所以代码中有一定的冗余(_0的else分支和_1完全一样)
  // 如果允许goto的话，也许性能会更好一些
  fn inorder_stk1(&self, f: &mut impl FnMut(u64)) {
    enum State<'a> {
      _0(&'a Tree),
      _1(&'a Tree),
    }
    let mut stk = vec![State::_0(self)];
    loop {
      match stk.pop() {
        Some(State::_0(t)) => if let Some(l) = &t.1 {
          stk.push(State::_1(t));
          stk.push(State::_0(l));
        } else {
          f(t.0);
          if let Some(r) = &t.2 { stk.push(State::_0(r)); }
        }
        Some(State::_1(t)) => {
          f(t.0);
          if let Some(r) = &t.2 { stk.push(State::_0(r)); }
        }
        None => return,
      }
    }
  }
}

树就随机插入数据来构建，这个代码就太平凡了，不展示了。随机插入1000000个数，求和20次取平均时间，结果如下：

inorder_rec inorder_gen inorder_stk inorder_stk1
36.85ms 212.3ms 36.45ms 44.55ms

可见现在的Generator的性能还是相当一般。标准的栈实现中序遍历的算法的性能和直接递归非常接近，甚至略优一点，不过对于现实中的问题不总是有这种优美的解法，但是都可以改写成基于栈+状态机的版本，这个性能虽然差一些，但也可以接受，我个人相比于用Generator更推荐它。