Rustでgeneratorを実現する

この記事のコードは https://github.com/gyu-don/minimal-generator に置いています。

# generator
最近のプログラミング言語には、generator記法がよく導入されています。generatorはイテレータのようなもので、関数と似た記法で値を `yield` できます。

Pythonでのgeneratorの例を示します。

```py:generatorの例
def one_two_three():
    yield 1
    yield 2
    yield 3

for i in one_two_three():
    print(i)
```

Rustでも、unstableな機能としてgeneratorは[あります](https://doc.rust-lang.org/std/ops/trait.Generator.html)が、まだ安定化されていません。
また、stableなRustでは、[genawaiter](https://docs.rs/genawaiter/latest/genawaiter/) crateを使うと、generatorが実現できます。

```rust:genawaiterを使ったgeneratorの例
fn main() {
    let one_two_three = gen!(|| {
        yield!(1);
        yield!(2);
        yield!(3);
    });
    
    for i in one_two_three.into_iter() {
        println!("{}", i);
    }
}
```

ところで、generatorって、一体どうやって実装できるのでしょうか。よくよく考えたら、厄介な話です。
generatorの値を作り出す側は、値を`yield`した際に、自分自身の状態を保存して、また、yieldした値を何らかの方法で、generatorの値を消費する側に受け渡さなければなりません。
また、generatorの値を消費する側は、`yield`された値を受け取り、また、実行が中断しているgeneratorを起こし、新しい値を要求しないといけません。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88256/013afc05-623c-c818-4883-b951aea4d9cc.png)

結構、これって面倒じゃないですか?

そこで、genawaiterの実装を読み解いて、generatorを実現してみました。
今回の実装は、中身を理解するための最小限のものなので、実際の利用は推奨しません。([`genawaiter`](https://docs.rs/genawaiter/latest/genawaiter/)を使ってください)

# やりたいこと概略 
フィボナッチ数列を返すgeneratorを作り、また、generatorから始めの10個、値を取り出す、という例を考えます。
やりたいことの概略をコードで示すと、以下のような感じです。

```rs:やりたいこと
fn main() {
    // yieldは、作り方が分からないので、panicする関数として作っておく
    let yield_ = |_| { unimplemented!() };
    // fibは、フィボナッチ数列を返すGeneratorとして定義したい
    let fib = || {
        let mut a = 1;
        let mut b = 1;
        loop {
            yield_(a);
            (a, b) = (b, a + b);
        }
    };
    // 先頭10個のfibの中身を取り出したい
    // let mut gen = fib()
    // for _ in 0..10 { ... }
}
```

# asyncなコルーチンを使って、中断を実現する
直感的に、generatorを実現する上で一番面倒なのは、generatorを関数のように書いたとき、`yield`で一旦関数を中断するところです。
これを自分で実現するのは非常に大変なのですが、非同期処理のための仕組みである`async fn`でgeneratorを書くことで、中断が実現します。

つまり、generatorから見ると、`yield`というのは、一旦値を上げて、また、再度呼ばれるまで`await`している、といえるわけです。

それを踏まえて、少し書き直してみましょう。

```rs:async/awaitを使った実装
async fn fib_generator(mut co: Co<u64>) {
    let mut a = 1u64;
    let mut b = 1u64;
    loop {
        co.yield_(a).await;
        (a, b) = (b, a + b);
    }
}

fn main() {
    // fibは、フィボナッチ数列を返すGeneratorとして定義したい
    // let fib = fib_generator(??);
    // 先頭10個のfibの中身を取り出したい
    // let mut gen = fib()
    // for _ in 0..10 { ... }
}
```

`co`という引数に、値を上げて、`await`するような実装になりました。
`co`には、`yield`された値を置く場所を作っておきます。また、`yield`すると`Future`が返ってくるはずです。とりあえず型だけ合わせた仮実装をしておきます。

```rs:Coを仮実装
use std::{rc::Rc, cell::Cell, future::Future};

pub struct Co<T> (Rc<Cell<Option<T>>>);

impl<T> Co<T> {
    pub fn yield_(&mut self, value: T) -> impl Future<Output=()> + '_ {
        DummyFuture()
    }
}

pub struct DummyFuture();
impl Future for DummyFuture {
    type Output = ();

    fn poll(self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Self::Output> {
        unimplemented!();
    }
}
```

# 実装をがんばる
`Co`の中身の型は、後に、値を受け取る側でも使うはずなので、別の型として定義しておきます。`Co::new`を定義します。

```rs:Co
type Airlock<T> = Rc<Cell<Option<T>>>;
pub struct Co<T> (Airlock<T>);

impl<T> Co<T> {
    pub fn new(airlock: Airlock<T>) -> Self {
        Self(airlock)
    }

    pub fn yield_(&mut self, value: T) -> impl Future<Output=()> + '_ {
        DummyFuture()
    }
}
```

また、generatorを`struct`として表すのですが、`yield`された中身を置く場所と、generatorを(`Future`として)持っておく場所が必要です。

```rs:Genを実装
pub struct Gen<T, F: Future> {
    airlock: Airlock<T>,
    future: Pin<Box<F>>,
}

impl<T, F: Future> Gen<T, F> {
    pub fn new(producer: impl FnOnce(Co<T>) -> F) -> Self {
        let airlock = Airlock::default();
        let future = Box::pin(producer(Co::new(airlock.clone())));
        Self { airlock, future }
    }
}
```

# `main`から見えるインタフェースを整える
generatorを再開するためのメソッドを`resume()`としましょう。これを呼び出したとき、`GeneratorState`型を返すことにします。`GeneratorState`は、以下のように定義し、`yield`された場合は`Yielded`、`yield`されずにgeneratorが終了した場合は`Complete`を返します。

```rs:GeneratorStateの定義
pub enum GeneratorState<T> {
    Yielded(T),
    Complete,
}
```

また、`resume()`は、`yield`された値を置く場所を空にして、次の値をもらいに行き、その値を`GeneratorState`として返します。

```rs:resumeの仮実装

pub struct Gen<T, F: Future> {
    airlock: Airlock<T>,
    future: Pin<Box<F>>,
}

impl<T, F: Future> Gen<T, F> {
    pub fn new(producer: impl FnOnce(Co<T>) -> F) -> Self {
        let airlock = Airlock::default();
        let future = Box::pin(producer(Co::new(airlock.clone())));
        Self { airlock, future }
    }

    pub fn resume(&mut self) -> GeneratorState<T> {
        self.airlock.replace(None);
        advance(self.future.as_mut(), &self.airlock)
    }
}

fn advance<T, F: Future>(future: Pin<&mut F>, airlock: &Airlock<T>) -> GeneratorState<T> {
    unimplemented!();
}
```

`main`関数は次のようになりました。

```rs:main
async fn fib_producer(mut co: Co<u64>) {
    let mut a = 1u64;
    let mut b = 1u64;
    loop {
        co.yield_(a).await;
        (a, b) = (b, a + b);
    }
}

fn main() {
    // fibは、フィボナッチ数列を返すGeneratorとして定義したい
    let mut fib = Gen::new(fib_producer);
    // 先頭10個のfibの中身を取り出したい
    for _ in 0..10 {
        match fib.resume() {
            GeneratorState::Yielded(n) => println!("{}", n),
            GeneratorState::Complete => {
                println!("Generator is completed.");
                break;
            }
        }
    }
}
```

# `yield_`をちゃんと実装する
仮実装だった`Future`をちゃんと実装します。
`yield_`は、上げられた値を、値の受け渡し場所にセットし、また、`Future`を返します。
その`Future`は、値が受け取られるのを待ちます。値が受け取られているかどうかは、受け渡し場所が空になっているかどうかで分かります。

```rs:yield_とFutureの実装
type Airlock<T> = Rc<Cell<Option<T>>>;
pub struct Co<T> (Airlock<T>);

impl<T> Co<T> {
    pub fn new(airlock: Airlock<T>) -> Self {
        Self(airlock)
    }

    pub fn yield_(&mut self, value: T) -> impl Future<Output=()> + '_ {
        self.0.replace(Some(value));
        Barrier(&self.0)
    }
}

pub struct Barrier<'a, T>(&'a Airlock<T>);
impl<'a, T> Future for Barrier<'a, T> {
    type Output = ();

    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
        match self.0.take() {
            Some(_) => Poll::Pending,
            None => Poll::Ready(()),
        }
    }
}
```

# 値を取ってくる側をちゃんと実装する
`advance`が仮実装だったので、ちゃんと実装します。`future`は、`poll`でつつくと「もういいよ」か「まだだよ」が返ってきます。「もういいよ」が返ってきたら、その`future`は終了なので、`GeneratorState::Complete`を返します。また、値がある場合は「まだだよ」を返ってきて、受け渡し場所には次の値が入れられている、という仕組みに実はなっています。それは、generatorが`await`するタイミングが、値を受け渡した後だからです。

```rs:advanceを実装
fn advance<T, F: Future>(future: Pin<&mut F>, airlock: &Airlock<T>) -> GeneratorState<T> {
    // futureからpollをするのに、Contextが必要。Contextを作るのにWakerが必要。Wakerはほぼ何もしないWakerを用意する。
    let waker = waker::create();
    let mut cx = Context::from_waker(&waker);

    match future.poll(&mut cx) {
        Poll::Ready(_) => {
            airlock.replace(None);
            GeneratorState::Complete
        },
        Poll::Pending => {
            match airlock.replace(None) {
                Some(v) => GeneratorState::Yielded(v),
                None => unreachable!(),
            }
        },
    }
}
```

また、`poll`をするために、`Waker`というものを作らないといけないのですが、この用途では、何もしない`Waker`を作って渡します。非常に関わりたくない雰囲気が出ていますが、怖がらずに読んでみたら、何もしていないことが分かります。

```rs:Wakerを実装
use std::{
    ptr,
    task::{RawWaker, RawWakerVTable, Waker},
};

pub fn create() -> Waker {
    // Safety: The waker points to a vtable with functions that do nothing. Doing
    // nothing is memory-safe.
    unsafe { Waker::from_raw(RAW_WAKER) }
}

const RAW_WAKER: RawWaker = RawWaker::new(ptr::null(), &VTABLE);
const VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop);

unsafe fn clone(_: *const ()) -> RawWaker {
    RAW_WAKER
}

unsafe fn wake(_: *const ()) {}

unsafe fn wake_by_ref(_: *const ()) {}

unsafe fn drop(_: *const ()) {}
```

これにより、generatorは完成で、以下のコードが動きます。

```rs:main
async fn fib_producer(mut co: Co<u64>) {
    let mut a = 1u64;
    let mut b = 1u64;
    loop {
        co.yield_(a).await;
        (a, b) = (b, a + b);
    }
}

fn main() {
    // 先頭10個のfibの中身を取り出したい
    for n in Gen::new(fib_producer).into_iter().take(10) {
        println!("{}", n);
    }
}
```

フィボナッチ数列が無事出力されています。

```動作結果
1
1
2
3
5
8
13
21
34
55
```

# generatorをイテレータにする

これで完成はしたのですが、`GeneratorState`を返すのでは使い勝手がよくないため、`IntoIterator`を実装します。

```rs:イテレータにする
impl<T, F: Future<Output=()>> IntoIterator for Gen<T, F> {
    type Item = T;
    type IntoIter = IntoIter<T, F>;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter(self)
    }
}

pub struct IntoIter<T, F> (Gen<T, F>);

impl<T, F: Future<Output=()>> Iterator for IntoIter<T, F> {
    type Item = T;

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

これにより、generatorをイテレータとして使うことができます。

```rs:main
async fn fib_producer(mut co: Co<u64>) {
    let mut a = 1u64;
    let mut b = 1u64;
    loop {
        co.yield_(a).await;
        (a, b) = (b, a + b);
    }
}

fn main() {
    // 先頭10個のfibの中身を取り出したい
    for n in Gen::new(fib_producer).into_iter().take(10) {
        println!("{}", n);
    }
```

# コード
[こちら](https://github.com/gyu-don/minimal-generator)に公開しています。
コミット履歴を追えば、この記事の好きな箇所の実装を見ることができます。