From 67689b3b4244ae83e119c0978e127686da487a2c Mon Sep 17 00:00:00 2001
From: SparrowLii <liyuan179@huawei.com>
Date: Sun, 29 Jan 2023 17:36:26 +0800
Subject: [PATCH 1/2] use RefLock in `GatedSpans` of Session

---
 compiler/rustc_data_structures/src/lib.rs  |  1 +
 compiler/rustc_data_structures/src/sync.rs | 88 ++++++++++++++++++++++
 compiler/rustc_session/src/parse.rs        |  4 +-
 3 files changed, 91 insertions(+), 2 deletions(-)

diff --git a/compiler/rustc_data_structures/src/lib.rs b/compiler/rustc_data_structures/src/lib.rs
index 954e84c303b83..7fab8954cb19f 100644
--- a/compiler/rustc_data_structures/src/lib.rs
+++ b/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)]
diff --git a/compiler/rustc_data_structures/src/sync.rs b/compiler/rustc_data_structures/src/sync.rs
index ed5341c40ef08..b912b7b6ff23b 100644
--- a/compiler/rustc_data_structures/src/sync.rs
+++ b/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>);
 
diff --git a/compiler/rustc_session/src/parse.rs b/compiler/rustc_session/src/parse.rs
index 2aa8ca9e4a919..117b1ae5ba167 100644
--- a/compiler/rustc_session/src/parse.rs
+++ b/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 {

From 6eded0d3d54bfd8bbb2416aa6682592cd03b9343 Mon Sep 17 00:00:00 2001
From: SparrowLii <liyuan179@huawei.com>
Date: Mon, 30 Jan 2023 10:17:52 +0800
Subject: [PATCH 2/2] disable thread check in non-parallel mode

---
 compiler/rustc_data_structures/src/sync.rs | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/compiler/rustc_data_structures/src/sync.rs b/compiler/rustc_data_structures/src/sync.rs
index b912b7b6ff23b..81628166fe6b3 100644
--- a/compiler/rustc_data_structures/src/sync.rs
+++ b/compiler/rustc_data_structures/src/sync.rs
@@ -488,7 +488,10 @@ impl<T> RefLock<T> {
 
     #[inline(always)]
     fn assert_thread(&self) {
+        #[cfg(parallel_compiler)]
         assert_eq!(get_thread_id(), self.thread_id);
+        #[cfg(not(parallel_compiler))]
+        return;
     }
 
     #[inline(always)]