この記事のコードは https://github.com/gyu-don/minimal-generator に置いています。
generator
最近のプログラミング言語には、generator記法がよく導入されています。generatorはイテレータのようなもので、関数と似た記法で値を yield できます。
Pythonでの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が実現できます。
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 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が返ってくるはずです。とりあえず型だけ合わせた仮実装をしておきます。
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として)持っておく場所が必要です。
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を返します。
pub enum GeneratorState<T> {
Yielded(T),
Complete,
}
また、resume()は、yieldされた値を置く場所を空にして、次の値をもらいに行き、その値をGeneratorStateとして返します。
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関数は次のようになりました。
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は、値が受け取られるのを待ちます。値が受け取られているかどうかは、受け渡し場所が空になっているかどうかで分かります。
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するタイミングが、値を受け渡した後だからです。
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を作って渡します。非常に関わりたくない雰囲気が出ていますが、怖がらずに読んでみたら、何もしていないことが分かります。
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は完成で、以下のコードが動きます。
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を実装します。
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をイテレータとして使うことができます。
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);
}
コード
こちらに公開しています。
コミット履歴を追えば、この記事の好きな箇所の実装を見ることができます。