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

generator

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

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

generatorの例

def one_two_three():
    yield 1
    yield 2
    yield 3

for i in one_two_three():
    print(i)

Rustでも、unstableな機能としてgeneratorはありますが、まだ安定化されていません。
また、stableなRustでは、genawaiter crateを使うと、generatorが実現できます。

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を起こし、新しい値を要求しないといけません。

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

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

やりたいこと概略

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

やりたいこと

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している、といえるわけです。

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

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が返ってくるはずです。とりあえず型だけ合わせた仮実装をしておきます。

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を定義します。

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として)持っておく場所が必要です。

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を返します。

GeneratorStateの定義

pub enum GeneratorState<T> {
    Yielded(T),
    Complete,
}

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

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関数は次のようになりました。

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は、値が受け取られるのを待ちます。値が受け取られているかどうかは、受け渡し場所が空になっているかどうかで分かります。

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するタイミングが、値を受け渡した後だからです。

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を作って渡します。非常に関わりたくない雰囲気が出ていますが、怖がらずに読んでみたら、何もしていないことが分かります。

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は完成で、以下のコードが動きます。

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

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

generatorをイテレータにする

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

イテレータにする

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をイテレータとして使うことができます。

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

コード

こちらに公開しています。
コミット履歴を追えば、この記事の好きな箇所の実装を見ることができます。

generator

やりたいこと概略

asyncなコルーチンを使って、中断を実現する

実装をがんばる

mainから見えるインタフェースを整える

yield_をちゃんと実装する

値を取ってくる側をちゃんと実装する

generatorをイテレータにする

コード

`main`から見えるインタフェースを整える

`yield_`をちゃんと実装する