Skip to content

Instantly share code, notes, and snippets.

@fegies
Created January 21, 2024 14:04
Show Gist options
  • Save fegies/5ac7b0171b161d77f877941c7353f0b3 to your computer and use it in GitHub Desktop.
Save fegies/5ac7b0171b161d77f877941c7353f0b3 to your computer and use it in GitHub Desktop.
use core::{
cell::Cell,
future::Future,
pin::{pin, Pin},
task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
};
pub struct ImpedanceMatcher<T> {
data: Cell<Option<T>>,
}
/// A simple future that will resolve after being polled once
pub struct PushFuture {
polled: bool,
}
impl Future for PushFuture {
type Output = ();
fn poll(
self: core::pin::Pin<&mut Self>,
_cx: &mut core::task::Context<'_>,
) -> core::task::Poll<Self::Output> {
if self.polled {
core::task::Poll::Ready(())
} else {
self.get_mut().polled = true;
Poll::Pending
}
}
}
impl<T> ImpedanceMatcher<T> {
pub fn push(&self, item: T) -> PushFuture {
self.data.set(Some(item));
PushFuture { polled: false }
}
pub fn new() -> Self {
Self {
data: Cell::new(None),
}
}
fn next(&self) -> Option<T> {
self.data.replace(None)
}
pub fn run<TFut, TRes>(
&self,
generator: TFut,
consumer: impl FnOnce(IteratorAdapter<'_, T, TFut>) -> TRes,
) -> TRes {
// infra
let waker = build_dummy_waker();
let context = Context::from_waker(&waker);
let future = pin!(generator);
let iter = IteratorAdapter {
context,
matcher: &self,
future,
};
consumer(iter)
}
}
pub struct IteratorAdapter<'a, T, F> {
matcher: &'a ImpedanceMatcher<T>,
future: Pin<&'a mut F>,
context: Context<'a>,
}
impl<T, F> Iterator for IteratorAdapter<'_, T, F>
where
F: Future,
{
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
let pinned = pin!(&mut self.future);
let _ = pinned.poll(&mut self.context);
self.matcher.next()
}
}
fn build_dummy_waker() -> Waker {
const VTABLE: RawWakerVTable = RawWakerVTable::new(
|_ptr| unreachable!(),
|_ptr| unreachable!(),
|_ptr| unreachable!(),
|_ptr| {},
);
unsafe {
let raw = RawWaker::new(core::ptr::null(), &VTABLE);
Waker::from_raw(raw)
}
}
#[cfg(test)]
#[test]
fn run_full_test() {
async fn generate_outputs(matcher: &ImpedanceMatcher<u32>) {
println!("gen_outputs");
matcher.push(0).await;
println!("after pushing 0");
matcher.push(1).await;
println!("after pushing 1");
matcher.push(2).await;
}
let matcher = ImpedanceMatcher::new();
let res = matcher.run(generate_outputs(&matcher), |iter| {
let mut res = Vec::new();
for item in iter {
println!("received {item}");
res.push(item);
}
res
});
assert_eq!([0, 1, 2].as_ref(), &res);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment