Compiler Explorer

Source code

//! An example using async "poll traits" which use manual "poll" functions
//! to create async functions.
//!
//! The Rust stdlib doesn't provide a way to run futures to completion,
//! so we have to define our own. To make the output easier to compare,
//! we create a waker which doesn't use any atomics.

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

use std::future::Future;
use std::pin::{Pin, pin};
use std::ptr;
use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};

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

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

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

pub trait Iterator {
    type Item;
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>>;
    #[inline(always)]
    async fn next(&mut self) -> Option<Self::Item>
    where
        Self: Unpin,
    {
        let fut = NextFuture { stream: self };
        fut.await
    }
}

pub struct NextFuture<'a, T: Unpin + ?Sized> {
    pub(crate) stream: &'a mut T,
}

impl<T: Iterator + Unpin + ?Sized> Future for NextFuture<'_, T> {
    type Output = Option<T::Item>;

fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        Pin::new(&mut *self.stream).poll_next(cx)
    }
}

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

#[derive(Clone, Debug)]
pub struct Once<T> {
    value: Option<T>,
}

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

fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<T>> {
        let this = unsafe { Pin::into_inner_unchecked(self) };
        Poll::Ready((&mut this.value).take())
    }
}

////////////////////////////////////////// Utilities:

/// Poll a non-Send future exactly once.
fn poll_once<T>(f: impl Future<Output = T>) -> Poll<T> {
    const WAKER: &Waker = {
        const RAW: RawWaker = {
            RawWaker::new(
                ptr::null(),
                &RawWakerVTable::new(no_clone, no_wake, no_wake, no_drop),
            )
        };
        fn no_clone(_: *const ()) -> RawWaker {
            RAW
        }
        fn no_wake(_: *const ()) {}
        fn no_drop(_: *const ()) {}
        &unsafe { Waker::from_raw(RAW) }
    };
    let mut ctx = Context::from_waker(&WAKER);
    pin!(f).poll(&mut ctx)
}