generators.rs

use std::future::Future;
use std::pin::Pin;
use std::task::{Poll, Context, Waker, RawWaker, RawWakerVTable};
use std::sync::mpsc::{sync_channel, Receiver};
use tokio::task::yield_now;

macro_rules! gen_iter {
    (
        @EMIT $yield_tx:ident
        []
        [$($output:tt)*]
    ) => {{
        $(
            //println!("{}", stringify!($output));
            $output
        )*
    }};
    // {}
    (
        @EMIT $yield_tx:ident
        [{$($toks:tt)*} $($rest:tt)*]
        [$($output:tt)*]
    ) => {
        gen_iter!(
            @EMIT $yield_tx
            [$($rest)*]
            [$($output)* {
                gen_iter!(
                    @EMIT $yield_tx
                    [$($toks)*]
                    []
                )
            }]
        )
    };
    // ()
    (
        @EMIT $yield_tx:ident
        [() $($rest:tt)*]
        [$($output:tt)*]
    ) => {
        gen_iter!(
            @EMIT $yield_tx
            [$($rest)*]
            [$($output)* ()]
        )
    };
    (
        @EMIT $yield_tx:ident
        [($($toks:tt)*) $($rest:tt)*]
        [$($output:tt)*]
    ) => {
        gen_iter!(
            @EMIT $yield_tx
            [$($rest)*]
            [$($output)* (
                gen_iter!(
                    @EMIT $yield_tx
                    [$($toks)*]
                    []
                )
            )]
        )
    };
    // []
    (
        @EMIT $yield_tx:ident
        [[$($toks:tt),*] $($rest:tt)*]
        [$($output:tt)*]
    ) => {
        gen_iter!(
            @EMIT $yield_tx
            [$($rest)*]
            [$($output)* [
                $(
                    gen_iter!(
                        @EMIT $yield_tx
                        [$toks]
                        []
                    )
                ),*
            ]]
        )
    };
    // forbid .await
    (
        @EMIT $yield_tx:ident
        [. await $($rest:tt)*]
        [$($output:tt)*]
    ) => {
        compile_error!("await is not allowed in gen_iter")
    };
    // inject yield
    (
        @EMIT $yield_tx:ident
        [yield $expr:expr; $($rest:tt)*]
        [$($output:tt)*]
    ) => {
        gen_iter!(
            @EMIT $yield_tx
            [$($rest)*]
            [$($output)*
                $yield_tx.send($expr).unwrap();
                yield_now().await;
            ]
        )
    };
    // any other token
    (
        @EMIT $yield_tx:ident
        [$tok:tt $($rest:tt)*]
        [$($output:tt)*]
    ) => {
        gen_iter!(
            @EMIT $yield_tx
            [$($rest)*]
            [$($output)* $tok]
        )
    };
    // entry point
    ($($toks:tt)*) => {{
        let (yield_tx, yield_rx) = sync_channel(1);
        
        GenIter {
            future: Some(Box::pin(async {
                let yield_tx = yield_tx;
                gen_iter!(
                    @EMIT yield_tx
                    [$($toks)*]
                    []
                )
            })),
            yield_rx,
        }
    }};
}

fn main() {
    let fib = gen_iter! {
        let mut a = 0;
        let mut b = 1;
        
        loop {
            yield Some(a);
            
            let next = a + b;
            a = b;
            b = next;
        }
    };
    
    for (i, n) in fib.enumerate().take(10) {
        println!("fib({}) = {}", i, n);
    }
}

struct GenIter<'a, R> {
    future: Option<Pin<Box<dyn Future<Output = Option<R>> + 'a>>>,
    yield_rx: Receiver<Option<R>>,
}

impl<'a, R> Iterator for GenIter<'a, R> {
    type Item = R;

    fn next(&mut self) -> Option<R> {
        let future = self.future.as_mut()?;
        let waker = RawWaker::new(&(), &DUMMY_WAKER_VTABLE);
        let waker = unsafe { Waker::from_raw(waker) };
        let mut context = Context::from_waker(&waker);
    
        match future.as_mut().poll(&mut context) {
            Poll::Ready(result) => {
                self.future = None;
                result
            },
            Poll::Pending => {
                let result = self.yield_rx.recv().unwrap();
                result
            }
        }
    }
}

const DUMMY_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
    |_| RawWaker::new(&(), &DUMMY_WAKER_VTABLE),
    |_| (),
    |_| (),
    |_| (),
);