Auto merge of #30834 - reem:rwlock-read-guard-map, r=alexcrichton

This is very useful when the RwLock is synchronizing access to a data
structure and you would like to return or store guards which contain
references to data inside the data structure instead of the data structure
itself.

2 files changed, 228 insertions, 57 deletions

diff --git a/src/libstd/sync/mutex.rs b/src/libstd/sync/mutex.rs
index fe9f0371abd..e83ebd10612 100644
--- a/src/libstd/sync/mutex.rs
+++ b/src/libstd/sync/mutex.rs
@@ -172,7 +172,7 @@ pub struct MutexGuard<'a, T: ?Sized + 'a> {
     // funny underscores due to how Deref/DerefMut currently work (they
     // disregard field privacy).
     __lock: &'a StaticMutex,
-    __data: &'a UnsafeCell<T>,
+    __data: &'a mut T,
     __poison: poison::Guard,
 }
 
@@ -211,8 +211,10 @@ impl<T: ?Sized> Mutex<T> {
     /// this call will return an error once the mutex is acquired.
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn lock(&self) -> LockResult<MutexGuard<T>> {
-        unsafe { self.inner.lock.lock() }
-        MutexGuard::new(&*self.inner, &self.data)
+        unsafe {
+            self.inner.lock.lock();
+            MutexGuard::new(&*self.inner, &self.data)
+        }
     }
 
     /// Attempts to acquire this lock.
@@ -230,10 +232,12 @@ impl<T: ?Sized> Mutex<T> {
     /// acquired.
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn try_lock(&self) -> TryLockResult<MutexGuard<T>> {
-        if unsafe { self.inner.lock.try_lock() } {
-            Ok(try!(MutexGuard::new(&*self.inner, &self.data)))
-        } else {
-            Err(TryLockError::WouldBlock)
+        unsafe {
+            if self.inner.lock.try_lock() {
+                Ok(try!(MutexGuard::new(&*self.inner, &self.data)))
+            } else {
+                Err(TryLockError::WouldBlock)
+            }
         }
     }
 
@@ -338,17 +342,21 @@ impl StaticMutex {
     /// Acquires this lock, see `Mutex::lock`
     #[inline]
     pub fn lock(&'static self) -> LockResult<MutexGuard<()>> {
-        unsafe { self.lock.lock() }
-        MutexGuard::new(self, &DUMMY.0)
+        unsafe {
+            self.lock.lock();
+            MutexGuard::new(self, &DUMMY.0)
+        }
     }
 
     /// Attempts to grab this lock, see `Mutex::try_lock`
     #[inline]
     pub fn try_lock(&'static self) -> TryLockResult<MutexGuard<()>> {
-        if unsafe { self.lock.try_lock() } {
-            Ok(try!(MutexGuard::new(self, &DUMMY.0)))
-        } else {
-            Err(TryLockError::WouldBlock)
+        unsafe {
+            if self.lock.try_lock() {
+                Ok(try!(MutexGuard::new(self, &DUMMY.0)))
+            } else {
+                Err(TryLockError::WouldBlock)
+            }
         }
     }
 
@@ -369,32 +377,72 @@ impl StaticMutex {
 
 impl<'mutex, T: ?Sized> MutexGuard<'mutex, T> {
 
-    fn new(lock: &'mutex StaticMutex, data: &'mutex UnsafeCell<T>)
+    unsafe fn new(lock: &'mutex StaticMutex, data: &'mutex UnsafeCell<T>)
            -> LockResult<MutexGuard<'mutex, T>> {
         poison::map_result(lock.poison.borrow(), |guard| {
             MutexGuard {
                 __lock: lock,
-                __data: data,
+                __data: &mut *data.get(),
                 __poison: guard,
             }
         })
     }
+
+    /// Transform this guard to hold a sub-borrow of the original data.
+    ///
+    /// Applies the supplied closure to the data, returning a new lock
+    /// guard referencing the borrow returned by the closure.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// # #![feature(guard_map)]
+    /// # use std::sync::{Mutex, MutexGuard};
+    /// let x = Mutex::new(vec![1, 2]);
+    ///
+    /// {
+    ///     let mut y = MutexGuard::map(x.lock().unwrap(), |v| &mut v[0]);
+    ///     *y = 3;
+    /// }
+    ///
+    /// assert_eq!(&*x.lock().unwrap(), &[3, 2]);
+    /// ```
+    #[unstable(feature = "guard_map",
+               reason = "recently added, needs RFC for stabilization",
+               issue = "27746")]
+    pub fn map<U: ?Sized, F>(this: Self, cb: F) -> MutexGuard<'mutex, U>
+        where F: FnOnce(&'mutex mut T) -> &'mutex mut U
+    {
+        // Compute the new data while still owning the original lock
+        // in order to correctly poison if the callback panics.
+        let data = unsafe { ptr::read(&this.__data) };
+        let new_data = cb(data);
+
+        // We don't want to unlock the lock by running the destructor of the
+        // original lock, so just read the fields we need and forget it.
+        let (poison, lock) = unsafe {
+            (ptr::read(&this.__poison), ptr::read(&this.__lock))
+        };
+        mem::forget(this);
+
+        MutexGuard {
+            __lock: lock,
+            __data: new_data,
+            __poison: poison
+        }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<'mutex, T: ?Sized> Deref for MutexGuard<'mutex, T> {
     type Target = T;
 
-    fn deref(&self) -> &T {
-        unsafe { &*self.__data.get() }
-    }
+    fn deref(&self) -> &T {self.__data }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<'mutex, T: ?Sized> DerefMut for MutexGuard<'mutex, T> {
-    fn deref_mut(&mut self) -> &mut T {
-        unsafe { &mut *self.__data.get() }
-    }
+    fn deref_mut(&mut self) -> &mut T { self.__data }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -421,7 +469,7 @@ mod tests {
     use prelude::v1::*;
 
     use sync::mpsc::channel;
-    use sync::{Arc, Mutex, StaticMutex, Condvar};
+    use sync::{Arc, Mutex, StaticMutex, Condvar, MutexGuard};
     use sync::atomic::{AtomicUsize, Ordering};
     use thread;
 
@@ -665,4 +713,19 @@ mod tests {
         let comp: &[i32] = &[4, 2, 5];
         assert_eq!(&*mutex.lock().unwrap(), comp);
     }
+
+    #[test]
+    fn test_mutex_guard_map_panic() {
+        let mutex = Arc::new(Mutex::new(vec![1, 2]));
+        let mutex2 = mutex.clone();
+
+        thread::spawn(move || {
+            let _ = MutexGuard::map::<usize, _>(mutex2.lock().unwrap(), |_| panic!());
+        }).join().unwrap_err();
+
+        match mutex.lock() {
+            Ok(r) => panic!("Lock on poisioned Mutex is Ok: {:?}", &*r),
+            Err(_) => {}
+        };
+    }
 }
diff --git a/src/libstd/sync/rwlock.rs b/src/libstd/sync/rwlock.rs
index 63ef7732ad6..1ecc8974369 100644
--- a/src/libstd/sync/rwlock.rs
+++ b/src/libstd/sync/rwlock.rs
@@ -121,7 +121,7 @@ pub const RW_LOCK_INIT: StaticRwLock = StaticRwLock::new();
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct RwLockReadGuard<'a, T: ?Sized + 'a> {
     __lock: &'a StaticRwLock,
-    __data: &'a UnsafeCell<T>,
+    __data: &'a T,
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -133,7 +133,7 @@ impl<'a, T: ?Sized> !marker::Send for RwLockReadGuard<'a, T> {}
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct RwLockWriteGuard<'a, T: ?Sized + 'a> {
     __lock: &'a StaticRwLock,
-    __data: &'a UnsafeCell<T>,
+    __data: &'a mut T,
     __poison: poison::Guard,
 }
 
@@ -177,8 +177,10 @@ impl<T: ?Sized> RwLock<T> {
     #[inline]
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn read(&self) -> LockResult<RwLockReadGuard<T>> {
-        unsafe { self.inner.lock.read() }
-        RwLockReadGuard::new(&*self.inner, &self.data)
+        unsafe {
+            self.inner.lock.read();
+            RwLockReadGuard::new(&*self.inner, &self.data)
+        }
     }
 
     /// Attempts to acquire this rwlock with shared read access.
@@ -201,10 +203,12 @@ impl<T: ?Sized> RwLock<T> {
     #[inline]
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn try_read(&self) -> TryLockResult<RwLockReadGuard<T>> {
-        if unsafe { self.inner.lock.try_read() } {
-            Ok(try!(RwLockReadGuard::new(&*self.inner, &self.data)))
-        } else {
-            Err(TryLockError::WouldBlock)
+        unsafe {
+            if self.inner.lock.try_read() {
+                Ok(try!(RwLockReadGuard::new(&*self.inner, &self.data)))
+            } else {
+                Err(TryLockError::WouldBlock)
+            }
         }
     }
 
@@ -225,8 +229,10 @@ impl<T: ?Sized> RwLock<T> {
     #[inline]
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn write(&self) -> LockResult<RwLockWriteGuard<T>> {
-        unsafe { self.inner.lock.write() }
-        RwLockWriteGuard::new(&*self.inner, &self.data)
+        unsafe {
+            self.inner.lock.write();
+            RwLockWriteGuard::new(&*self.inner, &self.data)
+        }
     }
 
     /// Attempts to lock this rwlock with exclusive write access.
@@ -249,10 +255,12 @@ impl<T: ?Sized> RwLock<T> {
     #[inline]
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn try_write(&self) -> TryLockResult<RwLockWriteGuard<T>> {
-        if unsafe { self.inner.lock.try_write() } {
-            Ok(try!(RwLockWriteGuard::new(&*self.inner, &self.data)))
-        } else {
-            Err(TryLockError::WouldBlock)
+        unsafe {
+            if self.inner.lock.try_write() {
+                Ok(try!(RwLockWriteGuard::new(&*self.inner, &self.data)))
+            } else {
+                Err(TryLockError::WouldBlock)
+            }
         }
     }
 
@@ -360,8 +368,10 @@ impl StaticRwLock {
     /// See `RwLock::read`.
     #[inline]
     pub fn read(&'static self) -> LockResult<RwLockReadGuard<'static, ()>> {
-        unsafe { self.lock.read() }
-        RwLockReadGuard::new(self, &DUMMY.0)
+        unsafe {
+            self.lock.read();
+            RwLockReadGuard::new(self, &DUMMY.0)
+        }
     }
 
     /// Attempts to acquire this lock with shared read access.
@@ -370,10 +380,12 @@ impl StaticRwLock {
     #[inline]
     pub fn try_read(&'static self)
                     -> TryLockResult<RwLockReadGuard<'static, ()>> {
-        if unsafe { self.lock.try_read() } {
-            Ok(try!(RwLockReadGuard::new(self, &DUMMY.0)))
-        } else {
-            Err(TryLockError::WouldBlock)
+        unsafe {
+            if self.lock.try_read(){
+                Ok(try!(RwLockReadGuard::new(self, &DUMMY.0)))
+            } else {
+                Err(TryLockError::WouldBlock)
+            }
         }
     }
 
@@ -383,8 +395,10 @@ impl StaticRwLock {
     /// See `RwLock::write`.
     #[inline]
     pub fn write(&'static self) -> LockResult<RwLockWriteGuard<'static, ()>> {
-        unsafe { self.lock.write() }
-        RwLockWriteGuard::new(self, &DUMMY.0)
+        unsafe {
+            self.lock.write();
+            RwLockWriteGuard::new(self, &DUMMY.0)
+        }
     }
 
     /// Attempts to lock this rwlock with exclusive write access.
@@ -393,10 +407,12 @@ impl StaticRwLock {
     #[inline]
     pub fn try_write(&'static self)
                      -> TryLockResult<RwLockWriteGuard<'static, ()>> {
-        if unsafe { self.lock.try_write() } {
-            Ok(try!(RwLockWriteGuard::new(self, &DUMMY.0)))
-        } else {
-            Err(TryLockError::WouldBlock)
+        unsafe {
+            if self.lock.try_write() {
+                Ok(try!(RwLockWriteGuard::new(self, &DUMMY.0)))
+            } else {
+                Err(TryLockError::WouldBlock)
+            }
         }
     }
 
@@ -412,48 +428,124 @@ impl StaticRwLock {
 }
 
 impl<'rwlock, T: ?Sized> RwLockReadGuard<'rwlock, T> {
-    fn new(lock: &'rwlock StaticRwLock, data: &'rwlock UnsafeCell<T>)
+    unsafe fn new(lock: &'rwlock StaticRwLock, data: &'rwlock UnsafeCell<T>)
            -> LockResult<RwLockReadGuard<'rwlock, T>> {
         poison::map_result(lock.poison.borrow(), |_| {
             RwLockReadGuard {
                 __lock: lock,
-                __data: data,
+                __data: &*data.get(),
             }
         })
     }
+
+    /// Transform this guard to hold a sub-borrow of the original data.
+    ///
+    /// Applies the supplied closure to the data, returning a new lock
+    /// guard referencing the borrow returned by the closure.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// # #![feature(guard_map)]
+    /// # use std::sync::{RwLockReadGuard, RwLock};
+    /// let x = RwLock::new(vec![1, 2]);
+    ///
+    /// let y = RwLockReadGuard::map(x.read().unwrap(), |v| &v[0]);
+    /// assert_eq!(*y, 1);
+    /// ```
+    #[unstable(feature = "guard_map",
+               reason = "recently added, needs RFC for stabilization",
+               issue = "27746")]
+    pub fn map<U: ?Sized, F>(this: Self, cb: F) -> RwLockReadGuard<'rwlock, U>
+        where F: FnOnce(&'rwlock T) -> &'rwlock U
+    {
+        let new = RwLockReadGuard {
+            __lock: this.__lock,
+            __data: cb(this.__data)
+        };
+
+        mem::forget(this);
+
+        new
+    }
 }
 
 impl<'rwlock, T: ?Sized> RwLockWriteGuard<'rwlock, T> {
-    fn new(lock: &'rwlock StaticRwLock, data: &'rwlock UnsafeCell<T>)
+    unsafe fn new(lock: &'rwlock StaticRwLock, data: &'rwlock UnsafeCell<T>)
            -> LockResult<RwLockWriteGuard<'rwlock, T>> {
         poison::map_result(lock.poison.borrow(), |guard| {
             RwLockWriteGuard {
                 __lock: lock,
-                __data: data,
+                __data: &mut *data.get(),
                 __poison: guard,
             }
         })
     }
+
+    /// Transform this guard to hold a sub-borrow of the original data.
+    ///
+    /// Applies the supplied closure to the data, returning a new lock
+    /// guard referencing the borrow returned by the closure.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// # #![feature(guard_map)]
+    /// # use std::sync::{RwLockWriteGuard, RwLock};
+    /// let x = RwLock::new(vec![1, 2]);
+    ///
+    /// {
+    ///     let mut y = RwLockWriteGuard::map(x.write().unwrap(), |v| &mut v[0]);
+    ///     assert_eq!(*y, 1);
+    ///
+    ///     *y = 10;
+    /// }
+    ///
+    /// assert_eq!(&**x.read().unwrap(), &[10, 2]);
+    /// ```
+    #[unstable(feature = "guard_map",
+               reason = "recently added, needs RFC for stabilization",
+               issue = "27746")]
+    pub fn map<U: ?Sized, F>(this: Self, cb: F) -> RwLockWriteGuard<'rwlock, U>
+        where F: FnOnce(&'rwlock mut T) -> &'rwlock mut U
+    {
+        // Compute the new data while still owning the original lock
+        // in order to correctly poison if the callback panics.
+        let data = unsafe { ptr::read(&this.__data) };
+        let new_data = cb(data);
+
+        // We don't want to unlock the lock by running the destructor of the
+        // original lock, so just read the fields we need and forget it.
+        let (poison, lock) = unsafe {
+            (ptr::read(&this.__poison), ptr::read(&this.__lock))
+        };
+        mem::forget(this);
+
+        RwLockWriteGuard {
+            __lock: lock,
+            __data: new_data,
+            __poison: poison
+        }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<'rwlock, T: ?Sized> Deref for RwLockReadGuard<'rwlock, T> {
     type Target = T;
 
-    fn deref(&self) -> &T { unsafe { &*self.__data.get() } }
+    fn deref(&self) -> &T { self.__data }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<'rwlock, T: ?Sized> Deref for RwLockWriteGuard<'rwlock, T> {
     type Target = T;
 
-    fn deref(&self) -> &T { unsafe { &*self.__data.get() } }
+    fn deref(&self) -> &T { self.__data }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<'rwlock, T: ?Sized> DerefMut for RwLockWriteGuard<'rwlock, T> {
-    fn deref_mut(&mut self) -> &mut T {
-        unsafe { &mut *self.__data.get() }
+    fn deref_mut(&mut self) -> &mut T { self.__data
     }
 }
 
@@ -481,7 +573,7 @@ mod tests {
     use rand::{self, Rng};
     use sync::mpsc::channel;
     use thread;
-    use sync::{Arc, RwLock, StaticRwLock, TryLockError};
+    use sync::{Arc, RwLock, StaticRwLock, TryLockError, RwLockWriteGuard};
     use sync::atomic::{AtomicUsize, Ordering};
 
     #[derive(Eq, PartialEq, Debug)]
@@ -729,4 +821,20 @@ mod tests {
             Ok(x) => panic!("get_mut of poisoned RwLock is Ok: {:?}", x),
         }
     }
+
+    #[test]
+    fn test_rwlock_write_map_poison() {
+        let rwlock = Arc::new(RwLock::new(vec![1, 2]));
+        let rwlock2 = rwlock.clone();
+
+        thread::spawn(move || {
+            let _ = RwLockWriteGuard::map::<usize, _>(rwlock2.write().unwrap(), |_| panic!());
+        }).join().unwrap_err();
+
+        match rwlock.read() {
+            Ok(r) => panic!("Read lock on poisioned RwLock is Ok: {:?}", &*r),
+            Err(_) => {}
+        };
+    }
 }
+