From 6047229bf262ed6995ceb462971e67e4bc7bd6f6 Mon Sep 17 00:00:00 2001
From: Wez Furlong <wez@wezfurlong.org>
Date: Wed, 14 Aug 2019 09:30:00 -0700
Subject: [PATCH] potentially make our promise/future crate await compatible

---
 promise/src/lib.rs | 33 ++++++++++++++++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

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<T> = Box<dyn FnOnce(Fallible<T>) + Send>;
 pub type SpawnFunc = Box<dyn FnOnce() + Send>;
@@ -59,6 +61,7 @@ enum PromiseState<T> {
 enum FutureState<T> {
     Waiting(Arc<Core<T>>),
     Ready(Result<T, Error>),
+    Resolved,
 }
 
 struct CoreData<T> {
@@ -208,6 +211,7 @@ impl<T: Send + 'static> Future<T> {
                     locked.propagate = Some(f);
                 }
             }
+            FutureState::Resolved => panic!("cannot chain a Resolved future"),
         }
     }
 
@@ -224,6 +228,7 @@ impl<T: Send + 'static> Future<T> {
                 }
             }
             FutureState::Ready(result) => result,
+            FutureState::Resolved => failure::bail!("Future is already Resolved"),
         }
     }
 
@@ -233,7 +238,7 @@ impl<T: Send + 'static> Future<T> {
                 let locked = core.data.lock().unwrap();
                 locked.result.is_some()
             }
-            FutureState::Ready(_) => true,
+            FutureState::Ready(_) | FutureState::Resolved => true,
         }
     }
 
@@ -319,6 +324,32 @@ impl<T: Send + 'static> Future<T> {
     }
 }
 
+impl<T: Send + 'static> std::future::Future for Future<T> {
+    type Output = Result<T, Error>;
+
+    fn poll(self: Pin<&mut Self>, _ctx: &mut Context) -> Poll<Self::Output> {
+        // 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::*;