Speedup future::poll_fn

Author: kpp

src/future.rs

@@ -1,6 +1,7 @@
 use futures::future::Future;
 use futures::stream::Stream;
 
+use core::pin::Pin;
 use core::task::{Context, Poll};
 
 /// Create a future that is immediately ready with a value.
@@ -502,18 +503,21 @@ pub fn poll_fn<F, T>(f: F) -> impl Future<Output = T>
 where
     F: FnMut(&mut Context) -> Poll<T>,
 {
-    use std::future::from_generator;
-    use std::future::get_task_context;
-
-    from_generator(|| {
-        let mut f = f;
-        loop {
-            match get_task_context(|context: &mut Context| f(context)) {
-                Poll::Pending => yield,
-                Poll::Ready(value) => return value,
-            }
+    struct PollFn<F> {
+        f: F,
+    }
+    impl<F> Unpin for PollFn<F> {}
+    impl<T, F> Future for PollFn<F>
+    where
+        F: FnMut(&mut Context<'_>) -> Poll<T>,
+    {
+        type Output = T;
+
+        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
+            (&mut self.f)(cx)
         }
-    })
+    }
+    PollFn { f }
 }
 
 #[cfg(test)]

src/stream.rs

@@ -30,13 +30,12 @@ use pin_utils::pin_mut;
 /// assert_eq!(next(&mut stream).await, None);
 /// # });
 /// ```
-pub async fn next<St>(stream: &mut St) -> Option<St::Item>
+pub fn next<'a, St>(mut stream: &'a mut St) -> impl Future<Output = Option<St::Item>> + 'a
 where
     St: Stream + Unpin,
 {
     use crate::future::poll_fn;
-    let future_next = poll_fn(|context| Pin::new(&mut *stream).poll_next(context));
-    future_next.await
+    poll_fn(move |context| Pin::new(&mut stream).poll_next(context))
 }
 
 /// Collect all of the values of this stream into a vector, returning a