Address review comments

Author: Diggsey

futures-util/src/future/future/map.rs

@@ -1,32 +1,52 @@
 use core::pin::Pin;
+use core::ptr;
 use futures_core::future::{FusedFuture, Future};
 use futures_core::task::{Context, Poll};
 use pin_project::{pin_project, project};
 
 use crate::fns::FnOnce1;
 
-/// Future for the [`map`](super::FutureExt::map) method.
+/// Internal Map future
 #[pin_project]
 #[derive(Debug)]
 #[must_use = "futures do nothing unless you `.await` or poll them"]
-pub struct Map<Fut, F> {
-    #[pin]
-    future: Fut,
-    f: Option<F>,
+pub enum Map<Fut, F> {
+    Incomplete {
+        #[pin]
+        future: Fut,
+        f: F,
+    },
+    Complete,
+}
+
+// Helper type to mark a `Map` as complete without running its destructor.
+struct UnsafeMarkAsComplete<Fut, F>(*mut Map<Fut, F>);
+
+impl<Fut, F> Drop for UnsafeMarkAsComplete<Fut, F> {
+    fn drop(&mut self) {
+        unsafe {
+            ptr::write(self.0, Map::Complete);
+        }
+    }
 }
 
 impl<Fut, F> Map<Fut, F> {
     /// Creates a new Map.
     pub(crate) fn new(future: Fut, f: F) -> Map<Fut, F> {
-        Map { future, f: Some(f) }
+        Map::Incomplete { future, f }
     }
 }
 
 impl<Fut, F, T> FusedFuture for Map<Fut, F>
     where Fut: Future,
           F: FnOnce1<Fut::Output, Output=T>,
 {
-    fn is_terminated(&self) -> bool { self.f.is_none() }
+    fn is_terminated(&self) -> bool {
+        match self {
+            Map::Incomplete { .. } => false,
+            Map::Complete => true,
+        }
+    }
 }
 
 impl<Fut, F, T> Future for Map<Fut, F>
@@ -36,13 +56,32 @@ impl<Fut, F, T> Future for Map<Fut, F>
     type Output = T;
 
     #[project]
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
-        #[project]
-        let Map { future, f } = self.project();
-        let output = ready!(future.poll(cx));
-        let f = f.take()
-            .expect("Map must not be polled after it returned `Poll::Ready`");
-
-        Poll::Ready(f.call_once(output))
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
+        unsafe {
+            // Store this pointer for later...
+            let self_ptr: *mut Self = self.as_mut().get_unchecked_mut();
+
+            #[project]
+            match self.project() {
+                Map::Incomplete { mut future, f } => {
+                    let output = ready!(future.as_mut().poll(cx));
+    
+                    // Here be dragons
+                    let f = {
+                        // The ordering here is important, the call to `drop_in_place` must be
+                        // last as it may panic. Other lines must not panic.
+                        let _cleanup = UnsafeMarkAsComplete(self_ptr);
+                        let f = ptr::read(f);
+                        ptr::drop_in_place(future.get_unchecked_mut());
+                        f
+                    };
+
+                    // Phew, everything is back to normal, and we should be in the
+                    // `Complete` state!
+                    Poll::Ready(f.call_once(output))
+                },
+                Map::Complete => panic!("Map must not be polled after it returned `Poll::Ready`"),
+            }
+        }
     }
 }

futures-util/src/future/future/mod.rs

@@ -42,8 +42,13 @@ delegate_all!(
 
 #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
 pub use fuse::Fuse;
-#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
-pub use map::Map;
+
+delegate_all!(
+    /// Future for the [`flatten`](super::FutureExt::flatten) method.
+    Map<Fut, F>(
+        map::Map<Fut, F>
+    ): Debug + Future + FusedFuture + Sink + Stream + FusedStream + New[|x: Fut, f: F| map::Map::new(x, f)]
+);
 
 delegate_all!(
     /// Stream for the [`into_stream`](FutureExt::into_stream) method.

futures-util/src/future/try_future/mod.rs

@@ -43,9 +43,10 @@ delegate_all!(
 
 delegate_all!(
     /// Future for the [`try_flatten_stream`](TryFutureExt::try_flatten_stream) method.
-    TryFlattenStream<Fut, St>(
-        try_flatten::TryFlatten<Fut, St>
+    TryFlattenStream<Fut>(
+        try_flatten::TryFlatten<Fut, Fut::Ok>
     ): Debug + Future + FusedFuture + Sink + Stream + FusedStream + New[|x: Fut| try_flatten::TryFlatten::new(x)]
+    where Fut: TryFuture
 );
 
 #[cfg(feature = "sink")]
@@ -526,7 +527,7 @@ pub trait TryFutureExt: TryFuture {
     /// assert_eq!(list, Ok(vec![17, 18, 19]));
     /// # });
     /// ```
-    fn try_flatten_stream(self) -> TryFlattenStream<Self, Self::Ok>
+    fn try_flatten_stream(self) -> TryFlattenStream<Self>
     where
         Self::Ok: TryStream<Error = Self::Error>,
         Self: Sized,

futures-util/src/future/try_future/try_flatten.rs

@@ -57,7 +57,7 @@ impl<Fut> Future for TryFlatten<Fut, Fut::Ok>
                     self.set(TryFlatten::Empty);
                     break output;
                 },
-                TryFlatten::Empty => return Poll::Pending,
+                TryFlatten::Empty => panic!("TryFlatten polled after completion"),
             }
         })
     }

futures-util/src/future/try_future/try_flatten_err.rs

@@ -54,7 +54,7 @@ impl<Fut> Future for TryFlattenErr<Fut, Fut::Error>
                     self.set(TryFlattenErr::Empty);
                     break output;
                 },
-                TryFlattenErr::Empty => return Poll::Pending,
+                TryFlattenErr::Empty => panic!("TryFlattenErr polled after completion"),
             }
         })
     }

futures-util/src/stream/stream/filter.rs

@@ -80,7 +80,9 @@ impl<St, Fut, F> Stream for Filter<St, Fut, F>
         let Filter { mut stream, f, mut pending_fut, pending_item } = self.project();
         Poll::Ready(loop {
             if let Some(fut) = pending_fut.as_mut().as_pin_mut() {
-                if ready!(fut.poll(cx)) {
+                let res = ready!(fut.poll(cx));
+                pending_fut.set(None);
+                if res {
                     break pending_item.take();
                 }
                 *pending_item = None;

futures-util/src/stream/stream/mod.rs

@@ -46,8 +46,14 @@ mod filter_map;
 pub use self::filter_map::FilterMap;
 
 mod flatten;
-#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
-pub use self::flatten::Flatten;
+
+delegate_all!(
+    /// Stream for the [`inspect`](StreamExt::inspect) method.
+    Flatten<St>(
+        flatten::Flatten<St, St::Item>
+    ): Debug + Sink + Stream + FusedStream + AccessInner[St, (.)] + New[|x: St| flatten::Flatten::new(x)]
+    where St: Stream
+);
 
 mod fold;
 #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
@@ -91,8 +97,8 @@ pub use self::map::Map;
 delegate_all!(
     /// Stream for the [`flat_map`](StreamExt::flat_map) method.
     FlatMap<St, U, F>(
-        Flatten<Map<St, F>, U>
-    ): Debug + Sink + Stream + FusedStream + AccessInner[St, (. .)] + New[|x: St, f: F| Flatten::new(Map::new(x, f))]
+        flatten::Flatten<Map<St, F>, U>
+    ): Debug + Sink + Stream + FusedStream + AccessInner[St, (. .)] + New[|x: St, f: F| flatten::Flatten::new(Map::new(x, f))]
 );
 
 mod next;
@@ -563,7 +569,7 @@ pub trait StreamExt: Stream {
     /// assert_eq!(output, vec![1, 2, 3, 4]);
     /// # });
     /// ```
-    fn flatten(self) -> Flatten<Self, Self::Item>
+    fn flatten(self) -> Flatten<Self>
     where
         Self::Item: Stream,
         Self: Sized,

futures-util/src/stream/try_stream/try_filter.rs

@@ -78,7 +78,9 @@ impl<St, Fut, F> Stream for TryFilter<St, Fut, F>
         let TryFilter { mut stream, f, mut pending_fut, pending_item } = self.project();
         Poll::Ready(loop {
             if let Some(fut) = pending_fut.as_mut().as_pin_mut() {
-                if ready!(fut.poll(cx)) {
+                let res = ready!(fut.poll(cx));
+                pending_fut.set(None);
+                if res {
                     break pending_item.take().map(Ok);
                 }
                 *pending_item = None;