From e0910be8fb85d6578d46be83e5fc4ff26fc874f1 Mon Sep 17 00:00:00 2001
From: "Abhishek C. Sharma" <abhishekcs@citrix.com>
Date: Sun, 3 Nov 2019 11:34:49 +0530
Subject: [PATCH] Added Future::flatten

---
 src/future/future.rs         | 27 +++++++++++++++++++++
 src/future/future/flatten.rs | 46 ++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)
 create mode 100644 src/future/future/flatten.rs

diff --git a/src/future/future.rs b/src/future/future.rs
index fe68517..1b88430 100644
--- a/src/future/future.rs
+++ b/src/future/future.rs
@@ -1,9 +1,12 @@
 cfg_unstable! {
     mod delay;
+    mod flatten;
 
     use std::time::Duration;
 
     use delay::DelayFuture;
+    use flatten::FlattenFuture;
+    use crate::future::IntoFuture;
 }
 
 extension_trait! {
@@ -129,6 +132,30 @@ extension_trait! {
         {
             DelayFuture::new(self, dur)
         }
+
+        /// Flatten out the execution of this future when the result itself
+        /// can be converted into another future.
+        ///
+        /// # Examples
+        ///
+        /// ```
+        /// # async_std::task::block_on(async {
+        /// use async_std::prelude::*;
+        ///
+        /// let nested_future = async { async { 1 } };
+        /// let future = nested_future.flatten();
+        /// assert_eq!(future.await, 1);
+        /// # })
+        /// ```
+        #[cfg_attr(feature = "docs", doc(cfg(unstable)))]
+        #[cfg(any(feature = "unstable", feature = "docs"))]
+        fn flatten(self) -> FlattenFuture<Self, <Self::Output as IntoFuture>::Future>
+        where
+            Self: Future + Sized,
+            Self::Output: IntoFuture
+        {
+           FlattenFuture::new(self)
+        }
     }
 
     impl<F: Future + Unpin + ?Sized> Future for Box<F> {
diff --git a/src/future/future/flatten.rs b/src/future/future/flatten.rs
new file mode 100644
index 0000000..27ee228
--- /dev/null
+++ b/src/future/future/flatten.rs
@@ -0,0 +1,46 @@
+use futures_core::ready;
+use std::pin::Pin;
+
+use crate::future::Future;
+use crate::future::IntoFuture;
+use crate::task::{Context, Poll};
+
+#[doc(hidden)]
+#[derive(Debug)]
+pub enum FlattenFuture<Fut1, Fut2> {
+    First(Fut1),
+    Second(Fut2),
+    Empty,
+}
+
+impl<Fut1, Fut2> FlattenFuture<Fut1, Fut2> {
+    pub fn new(future: Fut1) -> FlattenFuture<Fut1, Fut2> {
+        FlattenFuture::First(future)
+    }
+}
+
+impl<Fut1> Future for FlattenFuture<Fut1, <Fut1::Output as IntoFuture>::Future>
+where
+    Fut1: Future,
+    Fut1::Output: IntoFuture,
+{
+    type Output = <Fut1::Output as IntoFuture>::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = unsafe { self.get_unchecked_mut() };
+        loop {
+            match this {
+                FlattenFuture::First(fut1) => {
+                    let fut2 = ready!(unsafe { Pin::new_unchecked(fut1) }.poll(cx)).into_future();
+                    *this = FlattenFuture::Second(fut2);
+                }
+                FlattenFuture::Second(fut2) => {
+                    let v = ready!(unsafe { Pin::new_unchecked(fut2) }.poll(cx));
+                    *this = FlattenFuture::Empty;
+                    return Poll::Ready(v);
+                }
+                FlattenFuture::Empty => unreachable!(),
+            }
+        }
+    }
+}