diff --git a/promise/src/lib.rs b/promise/src/lib.rs index b043f1d84..ac0b74769 100644 --- a/promise/src/lib.rs +++ b/promise/src/lib.rs @@ -1,6 +1,8 @@ use failure::{Error, Fallible}; use failure_derive::*; +use std::pin::Pin; use std::sync::{Arc, Condvar, Mutex}; +use std::task::{Context, Poll}; type NextFunc = Box) + Send>; pub type SpawnFunc = Box; @@ -59,6 +61,7 @@ enum PromiseState { enum FutureState { Waiting(Arc>), Ready(Result), + Resolved, } struct CoreData { @@ -208,6 +211,7 @@ impl Future { locked.propagate = Some(f); } } + FutureState::Resolved => panic!("cannot chain a Resolved future"), } } @@ -224,6 +228,7 @@ impl Future { } } FutureState::Ready(result) => result, + FutureState::Resolved => failure::bail!("Future is already Resolved"), } } @@ -233,7 +238,7 @@ impl Future { let locked = core.data.lock().unwrap(); locked.result.is_some() } - FutureState::Ready(_) => true, + FutureState::Ready(_) | FutureState::Resolved => true, } } @@ -319,6 +324,32 @@ impl Future { } } +impl std::future::Future for Future { + type Output = Result; + + fn poll(self: Pin<&mut Self>, _ctx: &mut Context) -> Poll { + // This should be safe because we're not moving the Future, + // but instead replacing a field, and since no one is able to + // reference the state field, we should be ok with moving that. + let myself = unsafe { Pin::get_unchecked_mut(self) }; + + let state = std::mem::replace(&mut myself.state, FutureState::Resolved); + match state { + FutureState::Waiting(core) => { + let mut locked = core.data.lock().unwrap(); + if let Some(result) = locked.result.take() { + return Poll::Ready(result); + } + drop(locked); + myself.state = FutureState::Waiting(core); + Poll::Pending + } + FutureState::Ready(result) => Poll::Ready(result), + FutureState::Resolved => panic!("polling a Resolved Future"), + } + } +} + #[cfg(test)] mod test { use super::*;