From e8dc2c0571484b9410bdc0fe5b632ec45fc087c1 Mon Sep 17 00:00:00 2001 From: Taiki Endo 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 ef46d1a77..dc4c3a177 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 { - orig: S, - source: ManuallyDrop, +pin_project! { + /// A stream that will repeatedly yield the same list of elements. + #[derive(Debug)] + pub struct Cycle { + orig: S, + #[pin] + source: S, + } } impl Cycle @@ -18,15 +23,7 @@ where pub(crate) fn new(source: S) -> Self { Self { orig: source.clone(), - source: ManuallyDrop::new(source), - } - } -} - -impl Drop for Cycle { - 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> { - 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), } } }