Compiler Explorer

Source code

#![allow(incomplete_features)]
#![feature(async_fn_in_trait)]
#![feature(const_waker)]

use std::future::{self, Future};
use std::pin::pin;
use std::task::{Context, Poll, Wake};
use std::thread::{self, Thread};
use std::sync::Arc;

/////////////////////////////// Testing:

pub fn call_once() -> Poll<Option<usize>> {
    let mut iter = once(12usize);
    let output = poll_once(iter.next());
    output
}

/////////////////////// Implementation:

/// The async iterator trait
trait Iterator {
    type Item;
    async fn next(&mut self) -> Option<Self::Item>;
}

/// An async iterator which yields an item once
pub fn once<T>(t: T) -> Once<T> {
    Once { value: Some(t) }
}

pub struct Once<T> {
    value: Option<T>,
}

impl<T> Iterator for Once<T> {
    type Item = T;

async fn next(&mut self) -> Option<T> {
        future::poll_fn(|_cx| {
            Poll::Ready((&mut self.value).take())
        }).await
    }
}

////////////////////////////////////////// Utilities:
/// A waker that wakes up the current thread when called.
struct ThreadWaker(Thread);
impl Wake for ThreadWaker {
    fn wake(self: Arc<Self>) {
        self.0.unpark();
    }
}

/// Poll a non-Send future exactly once.
fn poll_once<T>(f: impl Future<Output = T>) -> Poll<T> {
    let t = thread::current();
    let waker = Arc::new(ThreadWaker(t)).into();
    let mut cx = Context::from_waker(&waker);
    pin!(f).poll(&mut cx)
}