Add guard map methods for transforming guards to contain sub-borrows.

This is very useful when the lock 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, 139 insertions, 5 deletions

diff --git a/src/libstd/sync/mutex.rs b/src/libstd/sync/mutex.rs
index 6b20e51967d..bdf8ca146fa 100644
--- a/src/libstd/sync/mutex.rs
+++ b/src/libstd/sync/mutex.rs
@@ -379,6 +379,43 @@ impl<'mutex, T: ?Sized> MutexGuard<'mutex, T> {
             }
         })
     }
+
+    /// 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.
+    ///
+    /// ```rust
+    /// # use std::sync::{Mutex, MutexGuard};
+    /// let x = Mutex::new(vec![1, 2]);
+    ///
+    /// {
+    ///     let y = MutexGuard::map(x.lock().unwrap(), |v| &mut v[0]);
+    ///     *y = 3;
+    /// }
+    ///
+    /// assert_eq!(&*x.lock(), &[3, 2]);
+    /// ```
+    #[unstable(feature = "guard_map",
+               reason = "recently added, needs RFC for stabilization",
+               issue = "0")]
+    pub fn map<U: ?Sized, F>(this: Self, cb: F) -> MutexGuard<'mutex, U>
+    where F: FnOnce(&'mutex mut T) -> &'mutex mut U {
+        let new_data = unsafe {
+            let data = cb(&mut *this.__data.get());
+            mem::transmute::<&'mutex mut U, &'mutex UnsafeCell<U>>(data)
+        };
+
+        let lock = unsafe { ptr::read(&this.__lock) };
+        let poison = unsafe { ptr::read(&this.__poison) };
+        mem::forget(this);
+
+        MutexGuard {
+            __lock: lock,
+            __data: new_data,
+            __poison: poison
+        }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -421,7 +458,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 +702,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 3dbef435481..04f00ed42d0 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")]
@@ -417,10 +417,37 @@ impl<'rwlock, T: ?Sized> RwLockReadGuard<'rwlock, T> {
         poison::map_result(lock.poison.borrow(), |_| {
             RwLockReadGuard {
                 __lock: lock,
-                __data: data,
+                __data: unsafe { &*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.
+    ///
+    /// ```rust
+    /// # 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 = "0")]
+    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> {
@@ -434,13 +461,52 @@ impl<'rwlock, T: ?Sized> RwLockWriteGuard<'rwlock, T> {
             }
         })
     }
+
+    /// 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.
+    ///
+    /// ```rust
+    /// # use std::sync::{RwLockWriteGuard, RwLock};
+    /// let x = RwLock::new(vec![1, 2]);
+    ///
+    /// {
+    ///     let 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 = "0")]
+    pub fn map<U: ?Sized, F>(this: Self, cb: F) -> RwLockWriteGuard<'rwlock, U>
+    where F: FnOnce(&'rwlock mut T) -> &'rwlock mut U {
+        let new_data = unsafe {
+            let data: &'rwlock mut T = &mut *this.__data.get();
+            mem::transmute::<&'rwlock mut U, &'rwlock UnsafeCell<U>>(cb(data))
+        };
+
+        let poison = unsafe { ptr::read(&this.__poison) };
+        let lock = unsafe { 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")]
@@ -481,7 +547,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 +795,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(_) => {}
+        };
+    }
 }
+