use RefLock in GatedSpans of Session

Author: SparrowLii

compiler/rustc_data_structures/src/lib.rs

@@ -11,6 +11,7 @@
 #![feature(associated_type_bounds)]
 #![feature(auto_traits)]
 #![feature(cell_leak)]
+#![feature(core_intrinsics)]
 #![feature(extend_one)]
 #![feature(hash_raw_entry)]
 #![feature(hasher_prefixfree_extras)]

compiler/rustc_data_structures/src/sync.rs

@@ -18,8 +18,11 @@
 //! depending on the value of cfg!(parallel_compiler).
 
 use crate::owning_ref::{Erased, OwningRef};
+use std::cell::{RefCell, RefMut};
 use std::collections::HashMap;
 use std::hash::{BuildHasher, Hash};
+use std::mem;
+use std::num::NonZeroUsize;
 use std::ops::{Deref, DerefMut};
 use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe};
 
@@ -460,6 +463,91 @@ impl<T: Clone> Clone for Lock<T> {
     }
 }
 
+fn next() -> NonZeroUsize {
+    static COUNTER: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(1);
+    NonZeroUsize::new(COUNTER.fetch_add(1, Ordering::SeqCst)).expect("more than usize::MAX threads")
+}
+
+pub(crate) fn get_thread_id() -> NonZeroUsize {
+    thread_local!(static THREAD_ID: NonZeroUsize = next());
+    THREAD_ID.with(|&x| x)
+}
+
+// `RefLock` is a thread-safe data structure because it can
+// only be used within the thread in which it was created.
+pub struct RefLock<T> {
+    val: RefCell<T>,
+    thread_id: NonZeroUsize,
+}
+
+impl<T> RefLock<T> {
+    #[inline(always)]
+    pub fn new(value: T) -> Self {
+        Self { val: RefCell::new(value), thread_id: get_thread_id() }
+    }
+
+    #[inline(always)]
+    fn assert_thread(&self) {
+        assert_eq!(get_thread_id(), self.thread_id);
+    }
+
+    #[inline(always)]
+    pub fn get_mut(&mut self) -> &mut T {
+        self.assert_thread();
+        self.val.get_mut()
+    }
+
+    #[inline(always)]
+    pub fn try_lock(&self) -> Option<RefMut<'_, T>> {
+        self.assert_thread();
+        self.val.try_borrow_mut().ok()
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    pub fn lock(&self) -> RefMut<'_, T> {
+        self.assert_thread();
+        self.val.borrow_mut()
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    pub fn with_lock<F: FnOnce(&mut T) -> R, R>(&self, f: F) -> R {
+        f(&mut *self.lock())
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    pub fn borrow(&self) -> RefMut<'_, T> {
+        self.lock()
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    pub fn borrow_mut(&self) -> RefMut<'_, T> {
+        self.lock()
+    }
+}
+
+unsafe impl<T> std::marker::Sync for RefLock<T> {}
+
+impl<T> Drop for RefLock<T> {
+    fn drop(&mut self) {
+        if mem::needs_drop::<T>() {
+            if get_thread_id() != self.thread_id {
+                panic!("destructor of fragile object ran on wrong thread");
+            }
+        }
+    }
+}
+
+impl<T: Default> Default for RefLock<T> {
+    #[inline]
+    fn default() -> Self {
+        RefLock::new(T::default())
+    }
+}
+
 #[derive(Debug, Default)]
 pub struct RwLock<T>(InnerRwLock<T>);
 

compiler/rustc_session/src/parse.rs

@@ -8,7 +8,7 @@ use crate::lint::{
 };
 use rustc_ast::node_id::NodeId;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexSet};
-use rustc_data_structures::sync::{Lock, Lrc};
+use rustc_data_structures::sync::{Lock, Lrc, RefLock};
 use rustc_errors::{emitter::SilentEmitter, ColorConfig, Handler};
 use rustc_errors::{
     fallback_fluent_bundle, Diagnostic, DiagnosticBuilder, DiagnosticId, DiagnosticMessage,
@@ -32,7 +32,7 @@ pub type CrateCheckConfig = CheckCfg<Symbol>;
 /// used and should be feature gated accordingly in `check_crate`.
 #[derive(Default)]
 pub struct GatedSpans {
-    pub spans: Lock<FxHashMap<Symbol, Vec<Span>>>,
+    pub spans: RefLock<FxHashMap<Symbol, Vec<Span>>>,
 }
 
 impl GatedSpans {