|
1 | | -use core::mem; |
| 1 | +use core::mem::Pin; |
2 | 2 |
|
3 | | -use futures_core::{Future, Poll, Async}; |
| 3 | +use futures_core::{Future, Poll}; |
4 | 4 | use futures_core::task; |
5 | 5 |
|
6 | 6 | #[must_use = "futures do nothing unless polled"] |
7 | 7 | #[derive(Debug)] |
8 | | -pub enum Chain<A, B, C> where A: Future { |
9 | | - First(A, C), |
10 | | - Second(B), |
11 | | - Done, |
| 8 | +pub enum Chain<Fut1, Fut2, Data> { |
| 9 | + First(Fut1, Option<Data>), |
| 10 | + Second(Fut2), |
12 | 11 | } |
13 | 12 |
|
14 | | -impl<A, B, C> Chain<A, B, C> |
15 | | - where A: Future, |
16 | | - B: Future, |
| 13 | +impl<Fut1, Fut2, Data> Chain<Fut1, Fut2, Data> |
| 14 | + where Fut1: Future, |
| 15 | + Fut2: Future, |
17 | 16 | { |
18 | | - pub fn new(a: A, c: C) -> Chain<A, B, C> { |
19 | | - Chain::First(a, c) |
| 17 | + pub fn new(fut1: Fut1, data: Data) -> Chain<Fut1, Fut2, Data> { |
| 18 | + Chain::First(fut1, Some(data)) |
20 | 19 | } |
21 | 20 |
|
22 | | - pub fn poll<F>(&mut self, cx: &mut task::Context, f: F) -> Poll<B::Item, B::Error> |
23 | | - where F: FnOnce(Result<A::Item, A::Error>, C) |
24 | | - -> Result<Result<B::Item, B>, B::Error>, |
| 21 | + pub fn poll<F>(mut self: Pin<Self>, cx: &mut task::Context, f: F) -> Poll<Fut2::Output> |
| 22 | + where F: FnOnce(Fut1::Output, Data) -> Fut2, |
25 | 23 | { |
26 | | - let a_result = match *self { |
27 | | - Chain::First(ref mut a, _) => { |
28 | | - match a.poll(cx) { |
29 | | - Ok(Async::Pending) => return Ok(Async::Pending), |
30 | | - Ok(Async::Ready(t)) => Ok(t), |
31 | | - Err(e) => Err(e), |
| 24 | + let mut f = Some(f); |
| 25 | + |
| 26 | + loop { |
| 27 | + // safe to `get_mut` here because we don't move out |
| 28 | + let fut2 = match *unsafe { Pin::get_mut(&mut self) } { |
| 29 | + Chain::First(ref mut fut1, ref mut data) => { |
| 30 | + // safe to create a new `Pin` because `fut1` will never move |
| 31 | + // before it's dropped. |
| 32 | + match unsafe { Pin::new_unchecked(fut1) }.poll(cx) { |
| 33 | + Poll::Pending => return Poll::Pending, |
| 34 | + Poll::Ready(t) => { |
| 35 | + (f.take().unwrap())(t, data.take().unwrap()) |
| 36 | + } |
| 37 | + } |
32 | 38 | } |
33 | | - } |
34 | | - Chain::Second(ref mut b) => return b.poll(cx), |
35 | | - Chain::Done => panic!("cannot poll a chained future twice"), |
36 | | - }; |
37 | | - let data = match mem::replace(self, Chain::Done) { |
38 | | - Chain::First(_, c) => c, |
39 | | - _ => panic!(), |
40 | | - }; |
41 | | - match f(a_result, data)? { |
42 | | - Ok(e) => Ok(Async::Ready(e)), |
43 | | - Err(mut b) => { |
44 | | - let ret = b.poll(cx); |
45 | | - *self = Chain::Second(b); |
46 | | - ret |
| 39 | + Chain::Second(ref mut fut2) => { |
| 40 | + // safe to create a new `Pin` because `fut2` will never move |
| 41 | + // before it's dropped; once we're in `Chain::Second` we stay |
| 42 | + // there forever. |
| 43 | + return unsafe { Pin::new_unchecked(fut2) }.poll(cx) |
| 44 | + } |
| 45 | + }; |
| 46 | + |
| 47 | + // safe because we're using the `&mut` to do an assignment, not for moving out |
| 48 | + unsafe { |
| 49 | + // note: it's safe to move the `fut2` here because we haven't yet polled it |
| 50 | + *Pin::get_mut(&mut self) = Chain::Second(fut2); |
47 | 51 | } |
48 | 52 | } |
49 | 53 | } |
|
0 commit comments