From e8dc2c0571484b9410bdc0fe5b632ec45fc087c1 Mon Sep 17 00:00:00 2001
From: Taiki Endo <te316e89@gmail.com>
Date: Mon, 2 Nov 2020 07:10:18 +0900
Subject: [PATCH] Fix double drop in StreamExt::cycle

---
 src/stream/stream/cycle.rs | 42 ++++++++++++++++----------------------
 1 file changed, 18 insertions(+), 24 deletions(-)

diff --git a/src/stream/stream/cycle.rs b/src/stream/stream/cycle.rs
index ef46d1a7..dc4c3a17 100644
--- a/src/stream/stream/cycle.rs
+++ b/src/stream/stream/cycle.rs
@@ -1,14 +1,19 @@
-use core::mem::ManuallyDrop;
 use core::pin::Pin;
 
+use futures_core::ready;
+use pin_project_lite::pin_project;
+
 use crate::stream::Stream;
 use crate::task::{Context, Poll};
 
-/// A stream that will repeatedly yield the same list of elements.
-#[derive(Debug)]
-pub struct Cycle<S> {
-    orig: S,
-    source: ManuallyDrop<S>,
+pin_project! {
+    /// A stream that will repeatedly yield the same list of elements.
+    #[derive(Debug)]
+    pub struct Cycle<S> {
+        orig: S,
+        #[pin]
+        source: S,
+    }
 }
 
 impl<S> Cycle<S>
@@ -18,15 +23,7 @@ where
     pub(crate) fn new(source: S) -> Self {
         Self {
             orig: source.clone(),
-            source: ManuallyDrop::new(source),
-        }
-    }
-}
-
-impl<S> Drop for Cycle<S> {
-    fn drop(&mut self) {
-        unsafe {
-            ManuallyDrop::drop(&mut self.source);
+            source,
         }
     }
 }
@@ -38,17 +35,14 @@ where
     type Item = S::Item;
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        unsafe {
-            let this = self.get_unchecked_mut();
+        let mut this = self.project();
 
-            match futures_core::ready!(Pin::new_unchecked(&mut *this.source).poll_next(cx)) {
-                Some(item) => Poll::Ready(Some(item)),
-                None => {
-                    ManuallyDrop::drop(&mut this.source);
-                    this.source = ManuallyDrop::new(this.orig.clone());
-                    Pin::new_unchecked(&mut *this.source).poll_next(cx)
-                }
+        match ready!(this.source.as_mut().poll_next(cx)) {
+            None => {
+                this.source.set(this.orig.clone());
+                this.source.poll_next(cx)
             }
+            item => Poll::Ready(item),
         }
     }
 }