From af1d2acd5c33137da8b74758f6c87675c44f9fa5 Mon Sep 17 00:00:00 2001
From: Vurich <jackefransham@gmail.com>
Date: Fri, 15 Dec 2017 10:45:06 +0100
Subject: [PATCH] Fix futures-mutex tests

---
 futures-mutex/src/lib.rs | 89 +++++++++++++++++++++++++++++++++++-----
 1 file changed, 78 insertions(+), 11 deletions(-)

diff --git a/futures-mutex/src/lib.rs b/futures-mutex/src/lib.rs
index 9ba4b57a..0aee079a 100644
--- a/futures-mutex/src/lib.rs
+++ b/futures-mutex/src/lib.rs
@@ -35,7 +35,7 @@
 //!     let future = AddTwo { lock: lock2 };
 //!
 //!     // This future will return the current value and the recovered lock.
-//!     let used_lock = lock1.lock().map(|b| (*b, b.unlock()));
+//!     let used_lock = lock1.into_lock().map(|b| (*b, b.unlock()));
 //!
 //!     let _ = future.join(used_lock).map(|(add_two, (value, _))| {
 //!         assert_eq!(add_two, value);
@@ -131,6 +131,17 @@ impl<T> Mutex<T> {
         }
     }
 
+    /// Convert this lock into a future that resolves to a guard that allows access to the data.
+    /// This function returns `MutexAcquire<T>`, which resolves to a `MutexGuard<T>`
+    /// guard type.
+    ///
+    /// The returned future will never return an error.
+    pub fn into_lock(self) -> MutexIntoAcquire<T> {
+        MutexIntoAcquire {
+            inner: self
+        }
+    }
+
     /// We unlock the mutex and wait for someone to lock. We try and unpark as many tasks as we
     /// can to prevents dead tasks from deadlocking the mutex, or tasks that have finished their
     /// critical section and were awakened.
@@ -200,6 +211,66 @@ impl<'a, T> Future for MutexAcquire<'a, T> {
     }
 }
 
+/// Future returned by `FutMutex::lock` which resolves to a guard when a lock is acquired.
+#[derive(Debug)]
+pub struct MutexIntoAcquire<T> {
+    inner: Mutex<T>
+}
+
+impl<T> Future for MutexIntoAcquire<T> {
+    type Item = MutexAcquired<T>;
+    type Error = ();
+
+    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        match self.inner.poll_lock() {
+            Async::Ready(r) => {
+                mem::forget(r);
+                Ok(MutexAcquired {
+                    inner: Mutex{ inner: self.inner.inner.clone() }
+                }.into())
+            },
+            Async::NotReady => Ok(Async::NotReady)
+        }
+    }
+}
+
+#[derive(Debug)]
+/// Resolved value of `FutMutexAcquire<T>` future
+///
+/// This value works like `FutMutexGuard<T>`, providing a RAII guard to the value `T` through
+/// `Deref` and `DerefMut`. Will unlock the lock when dropped; the original `FutMutex` can be
+/// recovered with `unlock()`.
+pub struct MutexAcquired<T> {
+    inner: Mutex<T>
+}
+
+impl<T> MutexAcquired<T> {
+    pub fn unlock(self) -> Mutex<T> {
+        Mutex {
+            inner: self.inner.inner.clone()
+        }
+    }
+}
+
+impl<T> Deref for MutexAcquired<T> {
+    type Target = T;
+    fn deref(&self) -> &Self::Target {
+        unsafe { &*self.inner.inner.data.get() }
+    }
+}
+
+impl<T> DerefMut for MutexAcquired<T> {
+    fn deref_mut(&mut self) -> &mut T {
+        unsafe { &mut *self.inner.inner.data.get() }
+    }
+}
+
+impl<T> Drop for MutexAcquired<T> {
+    fn drop(&mut self) {
+        self.inner.unlock();
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -208,14 +279,10 @@ mod tests {
     use futures::stream::{self, Stream};
     use std::thread;
 
-    pub fn unpark_noop() -> Arc<Notify> {
-        struct Foo;
+    struct Foo;
 
-        impl Notify for Foo {
-            fn notify(&self, id: usize) {}
-        }
-
-        Arc::new(Foo)
+    impl Notify for Foo {
+        fn notify(&self, id: usize) {}
     }
 
     #[test]
@@ -255,7 +322,7 @@ mod tests {
         });
 
         assert!(executor::spawn(future)
-                .poll_future_notify(unpark_noop())
+                .poll_future_notify(&Arc::new(Foo), 0)
                 .expect("failure in poll")
                 .is_ready());
     }
@@ -270,8 +337,8 @@ mod tests {
             a: Some(lock1),
             remaining: N,
         };
-        let b = stream::iter_ok((0..N).map(Ok::<_, ()>)).fold(lock2, |b, _n| {
-            b.lock().map(|mut b| {
+        let b = stream::iter_ok::<_, ()>(0..N).fold(lock2, |b, _n| {
+            b.into_lock().map(|mut b| {
                 *b += 1;
                 b.unlock()
             })