From a3e68704bc0f4f55a188e97b6c3d41cac50992a9 Mon Sep 17 00:00:00 2001
From: "Abhishek C. Sharma" <abhishekcs@citrix.com>
Date: Mon, 4 Nov 2019 13:58:14 +0530
Subject: [PATCH] Wrap state enum in public struct

---
 src/future/future/flatten.rs | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/future/future/flatten.rs b/src/future/future/flatten.rs
index 27ee2281..7b056142 100644
--- a/src/future/future/flatten.rs
+++ b/src/future/future/flatten.rs
@@ -5,9 +5,13 @@ use crate::future::Future;
 use crate::future::IntoFuture;
 use crate::task::{Context, Poll};
 
-#[doc(hidden)]
 #[derive(Debug)]
-pub enum FlattenFuture<Fut1, Fut2> {
+pub struct FlattenFuture<Fut1, Fut2> {
+    state: State<Fut1, Fut2>,
+}
+
+#[derive(Debug)]
+enum State<Fut1, Fut2> {
     First(Fut1),
     Second(Fut2),
     Empty,
@@ -15,7 +19,9 @@ pub enum FlattenFuture<Fut1, Fut2> {
 
 impl<Fut1, Fut2> FlattenFuture<Fut1, Fut2> {
     pub fn new(future: Fut1) -> FlattenFuture<Fut1, Fut2> {
-        FlattenFuture::First(future)
+        FlattenFuture {
+            state: State::First(future),
+        }
     }
 }
 
@@ -27,19 +33,19 @@ where
     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() };
+        let Self { state } = unsafe { self.get_unchecked_mut() };
         loop {
-            match this {
-                FlattenFuture::First(fut1) => {
+            match state {
+                State::First(fut1) => {
                     let fut2 = ready!(unsafe { Pin::new_unchecked(fut1) }.poll(cx)).into_future();
-                    *this = FlattenFuture::Second(fut2);
+                    *state = State::Second(fut2);
                 }
-                FlattenFuture::Second(fut2) => {
+                State::Second(fut2) => {
                     let v = ready!(unsafe { Pin::new_unchecked(fut2) }.poll(cx));
-                    *this = FlattenFuture::Empty;
+                    *state = State::Empty;
                     return Poll::Ready(v);
                 }
-                FlattenFuture::Empty => unreachable!(),
+                State::Empty => panic!("polled a completed future"),
             }
         }
     }