Fix build failure on latest nightly

Author: taiki-e

.github/workflows/ci.yml

@@ -26,6 +26,7 @@ jobs:
           #
           # Tests are not run as tests may require newer versions of Rust.
           - nightly-2019-08-21
+          - nightly-2019-12-09
           - nightly
     steps:
       - uses: actions/checkout@master

futures-async-stream-macro/src/visitor.rs

@@ -176,7 +176,7 @@ impl Visitor {
                 let mut __pinned = #base;
                 loop {
                     if let ::futures_async_stream::reexport::task::Poll::Ready(x) =
-                        ::futures_async_stream::stream::poll_with_tls_context(unsafe {
+                        ::futures_async_stream::future::poll_with_tls_context(unsafe {
                             ::futures_async_stream::reexport::pin::Pin::new_unchecked(&mut __pinned)
                         })
                     {

src/lib.rs

@@ -209,11 +209,10 @@
     no_crate_inject,
     attr(deny(warnings, rust_2018_idioms, single_use_lifetimes), allow(dead_code))
 ))]
-#![warn(unsafe_code)]
 #![warn(missing_docs)]
 #![warn(rust_2018_idioms, single_use_lifetimes, unreachable_pub)]
 #![warn(clippy::all)]
-#![feature(gen_future, generator_trait)]
+#![feature(generator_trait)]
 
 #[doc(inline)]
 pub use futures_async_stream_macro::for_await;
@@ -230,13 +229,289 @@ pub use futures_async_stream_macro::async_try_stream;
 #[doc(inline)]
 pub use futures_async_stream_macro::async_try_stream_block;
 
+use core::{cell::Cell, ptr::NonNull, task::Context};
+
+thread_local! {
+    static TLS_CX: Cell<Option<NonNull<Context<'static>>>> = Cell::new(None);
+}
+
+struct SetOnDrop(Option<NonNull<Context<'static>>>);
+
+impl Drop for SetOnDrop {
+    fn drop(&mut self) {
+        TLS_CX.with(|tls_cx| {
+            tls_cx.set(self.0.take());
+        });
+    }
+}
+
+// Safety: the returned guard must drop before `cx` is dropped and before
+// any previous guard is dropped.
+unsafe fn set_task_context(cx: &mut Context<'_>) -> SetOnDrop {
+    // transmute the context's lifetime to 'static so we can store it.
+    let cx = core::mem::transmute::<&mut Context<'_>, &mut Context<'static>>(cx);
+    let old_cx = TLS_CX.with(|tls_cx| tls_cx.replace(Some(NonNull::from(cx))));
+    SetOnDrop(old_cx)
+}
+
 // Not public API.
 #[doc(hidden)]
-pub mod stream;
+pub mod future {
+    use core::{
+        future::Future,
+        ops::{Generator, GeneratorState},
+        pin::Pin,
+        task::{Context, Poll},
+    };
+    use pin_project::pin_project;
+
+    use super::{set_task_context, SetOnDrop, TLS_CX};
+
+    // =================================================================================================
+    // GenFuture
+
+    /// Wrap a generator in a future.
+    ///
+    /// This function returns a `GenFuture` underneath, but hides it in `impl Trait` to give
+    /// better error messages (`impl Future` rather than `GenFuture<[closure.....]>`).
+    #[doc(hidden)]
+    pub fn from_generator<T: Generator<Yield = ()>>(x: T) -> impl Future<Output = T::Return> {
+        GenFuture(x)
+    }
+
+    /// A wrapper around generators used to implement `Future` for `async`/`await` code.
+    #[pin_project]
+    struct GenFuture<T>(#[pin] T);
+
+    #[doc(hidden)]
+    impl<T> Future for GenFuture<T>
+    where
+        T: Generator<Yield = ()>,
+    {
+        type Output = T::Return;
+
+        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+            let this = self.project();
+            let _guard = unsafe { set_task_context(cx) };
+            match this.0.resume() {
+                GeneratorState::Yielded(()) => Poll::Pending,
+                GeneratorState::Complete(x) => Poll::Ready(x),
+            }
+        }
+    }
+
+    // =================================================================================================
+    // Poll
+
+    /// Polls a future in the current thread-local task waker.
+    #[doc(hidden)]
+    pub fn poll_with_tls_context<F>(f: Pin<&mut F>) -> Poll<F::Output>
+    where
+        F: Future,
+    {
+        let cx_ptr = TLS_CX.with(|tls_cx| {
+            // Clear the entry so that nested `get_task_waker` calls
+            // will fail or set their own value.
+            tls_cx.replace(None)
+        });
+        let _reset = SetOnDrop(cx_ptr);
+
+        let mut cx_ptr = cx_ptr.expect("TLS Context not set.");
+
+        // Safety: we've ensured exclusive access to the context by
+        // removing the pointer from TLS, only to be replaced once
+        // we're done with it.
+        //
+        // The pointer that was inserted came from an `&mut Context<'_>`,
+        // so it is safe to treat as mutable.
+        unsafe { F::poll(f, cx_ptr.as_mut()) }
+    }
+}
+
+// Not public API.
+#[doc(hidden)]
+pub mod stream {
+    use core::{
+        future::Future,
+        marker::PhantomData,
+        ops::{Generator, GeneratorState},
+        pin::Pin,
+        task::{Context, Poll},
+    };
+    use futures_core::stream::Stream;
+    use pin_project::pin_project;
+
+    use super::{set_task_context, SetOnDrop, TLS_CX};
+
+    // =================================================================================================
+    // GenStream
+
+    /// Wrap a generator in a stream.
+    ///
+    /// This function returns a `GenStream` underneath, but hides it in `impl Trait` to give
+    /// better error messages (`impl Stream` rather than `GenStream<[closure.....]>`).
+    #[doc(hidden)]
+    pub fn from_generator<G, T>(gen: G) -> impl Stream<Item = T>
+    where
+        G: Generator<Yield = Poll<T>, Return = ()>,
+    {
+        GenStream { gen, _phantom: PhantomData }
+    }
+
+    /// A wrapper around generators used to implement `Stream` for `async`/`await` code.
+    #[pin_project]
+    struct GenStream<G, T> {
+        #[pin]
+        gen: G,
+        _phantom: PhantomData<T>,
+    }
+
+    impl<G, T> Stream for GenStream<G, T>
+    where
+        G: Generator<Yield = Poll<T>, Return = ()>,
+    {
+        type Item = T;
+
+        fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+            let this = self.project();
+            let _guard = unsafe { set_task_context(cx) };
+            match this.gen.resume() {
+                GeneratorState::Yielded(x) => x.map(Some),
+                GeneratorState::Complete(()) => Poll::Ready(None),
+            }
+        }
+    }
+
+    // =================================================================================================
+    // Poll
+
+    /// Polls a future in the current thread-local task waker.
+    #[doc(hidden)]
+    pub fn poll_next_with_tls_context<S>(s: Pin<&mut S>) -> Poll<Option<S::Item>>
+    where
+        S: Stream,
+    {
+        let cx_ptr = TLS_CX.with(|tls_cx| {
+            // Clear the entry so that nested `get_task_waker` calls
+            // will fail or set their own value.
+            tls_cx.replace(None)
+        });
+        let _reset = SetOnDrop(cx_ptr);
+
+        let mut cx_ptr = cx_ptr.expect("TLS Context not set.");
+
+        // Safety: we've ensured exclusive access to the context by
+        // removing the pointer from TLS, only to be replaced once
+        // we're done with it.
+        //
+        // The pointer that was inserted came from an `&mut Context<'_>`,
+        // so it is safe to treat as mutable.
+        unsafe { S::poll_next(s, cx_ptr.as_mut()) }
+    }
+
+    // =================================================================================================
+    // Next
+
+    // This is equivalent to the `futures::stream::StreamExt::next` method.
+    // But we want to make this crate dependency as small as possible, so we define our `next` function.
+    #[doc(hidden)]
+    pub fn next<S>(stream: &mut S) -> impl Future<Output = Option<S::Item>> + '_
+    where
+        S: Stream + Unpin,
+    {
+        Next { stream }
+    }
+
+    struct Next<'a, S> {
+        stream: &'a mut S,
+    }
+
+    impl<S> Future for Next<'_, S>
+    where
+        S: Stream + Unpin,
+    {
+        type Output = Option<S::Item>;
+
+        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+            Pin::new(&mut self.stream).poll_next(cx)
+        }
+    }
+}
 
 // Not public API.
 #[doc(hidden)]
-pub mod try_stream;
+pub mod try_stream {
+    use core::{
+        marker::PhantomData,
+        ops::{Generator, GeneratorState},
+        pin::Pin,
+        task::{Context, Poll},
+    };
+    use futures_core::stream::{FusedStream, Stream};
+    use pin_project::pin_project;
+
+    use super::set_task_context;
+
+    // =================================================================================================
+    // GenStream
+
+    /// Wrap a generator in a stream.
+    ///
+    /// This function returns a `GenStream` underneath, but hides it in `impl Trait` to give
+    /// better error messages (`impl Stream` rather than `GenStream<[closure.....]>`).
+    #[doc(hidden)]
+    pub fn from_generator<G, T, E>(
+        gen: G,
+    ) -> impl Stream<Item = Result<T, E>> + FusedStream<Item = Result<T, E>>
+    where
+        G: Generator<Yield = Poll<T>, Return = Result<(), E>>,
+    {
+        GenTryStream { gen, done: false, _phantom: PhantomData }
+    }
+
+    /// A wrapper around generators used to implement `Stream` for `async`/`await` code.
+    #[pin_project]
+    struct GenTryStream<G, T, E> {
+        #[pin]
+        gen: G,
+        done: bool,
+        _phantom: PhantomData<(T, E)>,
+    }
+
+    impl<G, T, E> Stream for GenTryStream<G, T, E>
+    where
+        G: Generator<Yield = Poll<T>, Return = Result<(), E>>,
+    {
+        type Item = Result<T, E>;
+
+        fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+            if self.done {
+                return Poll::Ready(None);
+            }
+
+            let this = self.project();
+            let _guard = unsafe { set_task_context(cx) };
+            let res = match this.gen.resume() {
+                GeneratorState::Yielded(x) => x.map(|x| Some(Ok(x))),
+                GeneratorState::Complete(Err(e)) => Poll::Ready(Some(Err(e))),
+                GeneratorState::Complete(Ok(())) => Poll::Ready(None),
+            };
+            if let Poll::Ready(Some(Err(_))) | Poll::Ready(None) = &res {
+                *this.done = true;
+            }
+            res
+        }
+    }
+
+    impl<G, T, E> FusedStream for GenTryStream<G, T, E>
+    where
+        G: Generator<Yield = Poll<T>, Return = Result<(), E>>,
+    {
+        fn is_terminated(&self) -> bool {
+            self.done
+        }
+    }
+}
 
 // Not public API.
 #[doc(hidden)]