Skip to content

Instantly share code, notes, and snippets.

@panicbit
Last active May 1, 2020 14:24
Show Gist options
  • Save panicbit/1033cb062495b32f7a1af025619e18aa to your computer and use it in GitHub Desktop.
Save panicbit/1033cb062495b32f7a1af025619e18aa to your computer and use it in GitHub Desktop.
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),
|_| (),
|_| (),
|_| (),
);
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment