impl Stream/Sink for FlattenSink/TryFlattenStream

futures-util/src/try_future/flatten_sink.rs

@@ -1,102 +1,76 @@
+use super::FlattenStreamSink;
 use core::pin::Pin;
 use futures_core::future::TryFuture;
+use futures_core::stream::{FusedStream, Stream, TryStream};
 use futures_core::task::{Context, Poll};
 use futures_sink::Sink;
-
-#[derive(Debug)]
-enum State<Fut, Si> {
-    Waiting(Fut),
-    Ready(Si),
-    Closed,
-}
-use self::State::*;
+use pin_utils::unsafe_pinned;
 
 /// Sink for the [`flatten_sink`](super::TryFutureExt::flatten_sink) method.
 #[derive(Debug)]
 #[must_use = "sinks do nothing unless polled"]
-pub struct FlattenSink<Fut, Si>(State<Fut, Si>);
-
-impl<Fut: Unpin, Si: Unpin> Unpin for FlattenSink<Fut, Si> {}
+pub struct FlattenSink<Fut, Si>
+where
+    Fut: TryFuture<Ok = Si>,
+{
+    inner: FlattenStreamSink<Fut>,
+}
 
 impl<Fut, Si> FlattenSink<Fut, Si>
 where
     Fut: TryFuture<Ok = Si>,
 {
-    pub(super) fn new(future: Fut) -> FlattenSink<Fut, Si> {
-        FlattenSink(Waiting(future))
-    }
+    unsafe_pinned!(inner: FlattenStreamSink<Fut>);
 
-    fn project_pin<'a>(
-        self: Pin<&'a mut Self>
-    ) -> State<Pin<&'a mut Fut>, Pin<&'a mut Si>> {
-        unsafe {
-            match &mut Pin::get_unchecked_mut(self).0 {
-                Waiting(f) => Waiting(Pin::new_unchecked(f)),
-                Ready(s) => Ready(Pin::new_unchecked(s)),
-                Closed => Closed,
-            }
+    pub(super) fn new(future: Fut) -> Self {
+        Self {
+            inner: FlattenStreamSink::new(future),
         }
     }
 }
 
+impl<Fut, S> FusedStream for FlattenSink<Fut, S>
+where
+    Fut: TryFuture<Ok = S>,
+    S: FusedStream,
+{
+    fn is_terminated(&self) -> bool {
+        self.inner.is_terminated()
+    }
+}
+
+impl<Fut, S> Stream for FlattenSink<Fut, S>
+where
+    Fut: TryFuture<Ok = S>,
+    S: TryStream<Error = Fut::Error>,
+{
+    type Item = Result<S::Ok, Fut::Error>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.inner().poll_next(cx)
+    }
+}
+
 impl<Fut, Si, Item> Sink<Item> for FlattenSink<Fut, Si>
 where
     Fut: TryFuture<Ok = Si>,
     Si: Sink<Item, SinkError = Fut::Error>,
 {
-    type SinkError = Si::SinkError;
+    type SinkError = Fut::Error;
 
-    fn poll_ready(
-        mut self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<Result<(), Self::SinkError>> {
-        let resolved_stream = match self.as_mut().project_pin() {
-            Ready(s) => return s.poll_ready(cx),
-            Waiting(f) => ready!(f.try_poll(cx))?,
-            Closed => panic!("poll_ready called after eof"),
-        };
-        self.set(FlattenSink(Ready(resolved_stream)));
-        if let Ready(resolved_stream) = self.project_pin() {
-            resolved_stream.poll_ready(cx)
-        } else {
-            unreachable!()
-        }
+    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::SinkError>> {
+        self.inner().poll_ready(cx)
     }
 
-    fn start_send(
-        self: Pin<&mut Self>,
-        item: Item,
-    ) -> Result<(), Self::SinkError> {
-        match self.project_pin() {
-            Ready(s) => s.start_send(item),
-            Waiting(_) => panic!("poll_ready not called first"),
-            Closed => panic!("start_send called after eof"),
-        }
+    fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::SinkError> {
+        self.inner().start_send(item)
     }
 
-    fn poll_flush(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<Result<(), Self::SinkError>> {
-        match self.project_pin() {
-            Ready(s) => s.poll_flush(cx),
-            // if sink not yet resolved, nothing written ==> everything flushed
-            Waiting(_) => Poll::Ready(Ok(())),
-            Closed => panic!("poll_flush called after eof"),
-        }
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::SinkError>> {
+        self.inner().poll_flush(cx)
     }
 
-    fn poll_close(
-        mut self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<Result<(), Self::SinkError>> {
-        let res = match self.as_mut().project_pin() {
-            Ready(s) => s.poll_close(cx),
-            Waiting(_) | Closed => Poll::Ready(Ok(())),
-        };
-        if res.is_ready() {
-            self.set(FlattenSink(Closed));
-        }
-        res
+    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::SinkError>> {
+        self.inner().poll_close(cx)
     }
 }

futures-util/src/try_future/flatten_stream_sink.rs

@@ -0,0 +1,180 @@
+use core::fmt;
+use core::pin::Pin;
+use futures_core::future::TryFuture;
+use futures_core::stream::{FusedStream, Stream, TryStream};
+use futures_core::task::{Context, Poll};
+use futures_sink::Sink;
+use pin_utils::unsafe_pinned;
+
+/// Stream for the [`flatten_stream_sink`](super::TryFutureExt::flatten_stream_sink) method.
+#[must_use = "streams do nothing unless polled"]
+pub(crate) struct FlattenStreamSink<Fut>
+where
+    Fut: TryFuture,
+{
+    state: State<Fut, Fut::Ok>,
+}
+
+impl<Fut> Unpin for FlattenStreamSink<Fut>
+where
+    Fut: TryFuture + Unpin,
+    Fut::Ok: Unpin,
+{
+}
+
+impl<Fut> fmt::Debug for FlattenStreamSink<Fut>
+where
+    Fut: TryFuture + fmt::Debug,
+    Fut::Ok: fmt::Debug,
+{
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt.debug_struct("FlattenStreamSink")
+            .field("state", &self.state)
+            .finish()
+    }
+}
+
+impl<Fut> FlattenStreamSink<Fut>
+where
+    Fut: TryFuture,
+{
+    unsafe_pinned!(state: State<Fut, Fut::Ok>);
+
+    pub(crate) fn new(future: Fut) -> Self {
+        Self {
+            state: State::Future(future),
+        }
+    }
+}
+
+#[derive(Debug)]
+enum State<Fut, S> {
+    // future is not yet called or called and not ready
+    Future(Fut),
+    // future resolved to Stream or Sink
+    StreamOrSink(S),
+    // future resolved to error
+    Done,
+}
+
+impl<Fut, S> State<Fut, S> {
+    fn get_pin_mut<'a>(self: Pin<&'a mut Self>) -> State<Pin<&'a mut Fut>, Pin<&'a mut S>> {
+        // safety: data is never moved via the resulting &mut reference
+        match unsafe { Pin::get_unchecked_mut(self) } {
+            // safety: the future we're re-pinning here will never be moved;
+            // it will just be polled, then dropped in place
+            State::Future(f) => State::Future(unsafe { Pin::new_unchecked(f) }),
+            // safety: the stream we're repinning here will never be moved;
+            // it will just be polled, then dropped in place
+            State::StreamOrSink(s) => State::StreamOrSink(unsafe { Pin::new_unchecked(s) }),
+            State::Done => State::Done,
+        }
+    }
+}
+
+impl<Fut> State<Fut, Fut::Ok>
+where
+    Fut: TryFuture,
+{
+    fn poll_future(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Fut::Error>> {
+        if let State::Future(f) = self.as_mut().get_pin_mut() {
+            match ready!(f.try_poll(cx)) {
+                Ok(s) => {
+                    // Future resolved to stream.
+                    // We do not return, but poll that
+                    // stream in the next loop iteration.
+                    self.set(State::StreamOrSink(s));
+                }
+                Err(e) => {
+                    // Future resolved to error.
+                    // We have neither a pollable stream nor a future.
+                    self.set(State::Done);
+                    return Poll::Ready(Err(e));
+                }
+            }
+        }
+        Poll::Ready(Ok(()))
+    }
+}
+
+impl<Fut> FusedStream for FlattenStreamSink<Fut>
+where
+    Fut: TryFuture,
+    Fut::Ok: FusedStream,
+{
+    fn is_terminated(&self) -> bool {
+        match &self.state {
+            State::Future(_) => false,
+            State::StreamOrSink(stream) => stream.is_terminated(),
+            State::Done => true,
+        }
+    }
+}
+
+impl<Fut> Stream for FlattenStreamSink<Fut>
+where
+    Fut: TryFuture,
+    Fut::Ok: TryStream<Error = Fut::Error>,
+{
+    type Item = Result<<Fut::Ok as TryStream>::Ok, Fut::Error>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        ready!(self.as_mut().state().poll_future(cx)?);
+        match self.as_mut().state().get_pin_mut() {
+            State::StreamOrSink(s) => s.try_poll_next(cx),
+            State::Done => Poll::Ready(None),
+            State::Future(_) => unreachable!(),
+        }
+    }
+}
+
+impl<Fut, Item> Sink<Item> for FlattenStreamSink<Fut>
+where
+    Fut: TryFuture,
+    Fut::Ok: Sink<Item, SinkError = Fut::Error>,
+{
+    type SinkError = Fut::Error;
+
+    fn poll_ready(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(), Self::SinkError>> {
+        ready!(self.as_mut().state().poll_future(cx)?);
+        match self.as_mut().state().get_pin_mut() {
+            State::StreamOrSink(s) => s.poll_ready(cx),
+            State::Done => panic!("poll_ready called after eof"),
+            State::Future(_) => unreachable!(),
+        }
+    }
+
+    fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::SinkError> {
+        match self.state().get_pin_mut() {
+            State::StreamOrSink(s) => s.start_send(item),
+            State::Future(_) => panic!("poll_ready not called first"),
+            State::Done => panic!("start_send called after eof"),
+        }
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::SinkError>> {
+        match self.state().get_pin_mut() {
+            State::StreamOrSink(s) => s.poll_flush(cx),
+            // if sink not yet resolved, nothing written ==> everything flushed
+            State::Future(_) => Poll::Ready(Ok(())),
+            State::Done => panic!("poll_flush called after eof"),
+        }
+    }
+
+    fn poll_close(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(), Self::SinkError>> {
+        let res = match self.as_mut().state().get_pin_mut() {
+            State::StreamOrSink(s) => s.poll_close(cx),
+            State::Future(_) | State::Done => Poll::Ready(Ok(())),
+        };
+        if res.is_ready() {
+            self.as_mut().state().set(State::Done);
+        }
+        res
+    }
+}

futures-util/src/try_future/mod.rs

@@ -68,6 +68,9 @@ pub use self::unwrap_or_else::UnwrapOrElse;
 mod try_chain;
 pub(crate) use self::try_chain::{TryChain, TryChainAction};
 
+mod flatten_stream_sink;
+pub(crate) use self::flatten_stream_sink::FlattenStreamSink;
+
 impl<Fut: ?Sized + TryFuture> TryFutureExt for Fut {}
 
 /// Adapters specific to [`Result`]-returning futures