From 2aa19feeb988710e6e9ca9e1c8a77f99e3fe7213 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= <john.kare.alsaker@gmail.com>
Date: Thu, 8 Mar 2018 04:50:43 +0100
Subject: [PATCH] Add with_lock, with_read_lock and with_write_lock

---
 src/librustc_data_structures/sync.rs | 126 ++++++++++++++++++---------
 1 file changed, 87 insertions(+), 39 deletions(-)

diff --git a/src/librustc_data_structures/sync.rs b/src/librustc_data_structures/sync.rs
index b1ab4eaa069..f17ab264bff 100644
--- a/src/librustc_data_structures/sync.rs
+++ b/src/librustc_data_structures/sync.rs
@@ -62,7 +62,7 @@ cfg_if! {
         pub use std::cell::RefMut as WriteGuard;
         pub use std::cell::RefMut as LockGuard;
 
-        pub use std::cell::RefCell as RwLock;
+        use std::cell::RefCell as InnerRwLock;
         use std::cell::RefCell as InnerLock;
 
         use std::cell::Cell;
@@ -159,13 +159,12 @@ cfg_if! {
 
         pub use parking_lot::MutexGuard as LockGuard;
 
-        use parking_lot;
-
         pub use std::sync::Arc as Lrc;
 
         pub use self::Lock as MTLock;
 
         use parking_lot::Mutex as InnerLock;
+        use parking_lot::RwLock as InnerRwLock;
 
         pub type MetadataRef = OwningRef<Box<Erased + Send + Sync>, [u8]>;
 
@@ -222,42 +221,6 @@ cfg_if! {
                 self.0.lock().take()
             }
         }
-
-        #[derive(Debug)]
-        pub struct RwLock<T>(parking_lot::RwLock<T>);
-
-        impl<T> RwLock<T> {
-            #[inline(always)]
-            pub fn new(inner: T) -> Self {
-                RwLock(parking_lot::RwLock::new(inner))
-            }
-
-            #[inline(always)]
-            pub fn borrow(&self) -> ReadGuard<T> {
-                if ERROR_CHECKING {
-                    self.0.try_read().expect("lock was already held")
-                } else {
-                    self.0.read()
-                }
-            }
-
-            #[inline(always)]
-            pub fn borrow_mut(&self) -> WriteGuard<T> {
-                if ERROR_CHECKING {
-                    self.0.try_write().expect("lock was already held")
-                } else {
-                    self.0.write()
-                }
-            }
-        }
-
-        // FIXME: Probably a bad idea
-        impl<T: Clone> Clone for RwLock<T> {
-            #[inline]
-            fn clone(&self) -> Self {
-                RwLock::new(self.borrow().clone())
-            }
-        }
     }
 }
 
@@ -383,6 +346,11 @@ impl<T> Lock<T> {
         self.0.borrow_mut()
     }
 
+    #[inline(always)]
+    pub fn with_lock<F: FnOnce(&mut T) -> R, R>(&self, f: F) -> R {
+        f(&mut *self.lock())
+    }
+
     #[inline(always)]
     pub fn borrow(&self) -> LockGuard<T> {
         self.lock()
@@ -401,3 +369,83 @@ impl<T: Clone> Clone for Lock<T> {
         Lock::new(self.borrow().clone())
     }
 }
+
+#[derive(Debug)]
+pub struct RwLock<T>(InnerRwLock<T>);
+
+impl<T> RwLock<T> {
+    #[inline(always)]
+    pub fn new(inner: T) -> Self {
+        RwLock(InnerRwLock::new(inner))
+    }
+
+    #[inline(always)]
+    pub fn into_inner(self) -> T {
+        self.0.into_inner()
+    }
+
+    #[inline(always)]
+    pub fn get_mut(&mut self) -> &mut T {
+        self.0.get_mut()
+    }
+
+    #[cfg(not(parallel_queries))]
+    #[inline(always)]
+    pub fn read(&self) -> ReadGuard<T> {
+        self.0.borrow()
+    }
+
+    #[cfg(parallel_queries)]
+    #[inline(always)]
+    pub fn read(&self) -> ReadGuard<T> {
+        if ERROR_CHECKING {
+            self.0.try_read().expect("lock was already held")
+        } else {
+            self.0.read()
+        }
+    }
+
+    #[inline(always)]
+    pub fn with_read_lock<F: FnOnce(&T) -> R, R>(&self, f: F) -> R {
+        f(&*self.read())
+    }
+
+    #[cfg(not(parallel_queries))]
+    #[inline(always)]
+    pub fn write(&self) -> WriteGuard<T> {
+        self.0.borrow_mut()
+    }
+
+    #[cfg(parallel_queries)]
+    #[inline(always)]
+    pub fn write(&self) -> WriteGuard<T> {
+        if ERROR_CHECKING {
+            self.0.try_write().expect("lock was already held")
+        } else {
+            self.0.write()
+        }
+    }
+
+    #[inline(always)]
+    pub fn with_write_lock<F: FnOnce(&mut T) -> R, R>(&self, f: F) -> R {
+        f(&mut *self.write())
+    }
+
+    #[inline(always)]
+    pub fn borrow(&self) -> ReadGuard<T> {
+        self.read()
+    }
+
+    #[inline(always)]
+    pub fn borrow_mut(&self) -> WriteGuard<T> {
+        self.write()
+    }
+}
+
+// FIXME: Probably a bad idea
+impl<T: Clone> Clone for RwLock<T> {
+    #[inline]
+    fn clone(&self) -> Self {
+        RwLock::new(self.borrow().clone())
+    }
+}