Try AtomicPtrs

Author: aerooneqq

compiler/rustc_middle/src/dep_graph/graph.rs

@@ -2,7 +2,7 @@ use std::assert_matches;
 use std::fmt::Debug;
 use std::hash::Hash;
 use std::sync::Arc;
-use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
+use std::sync::atomic::{AtomicPtr, AtomicU32, Ordering};
 
 use rustc_data_structures::fingerprint::{Fingerprint, PackedFingerprint};
 use rustc_data_structures::fx::FxHashSet;
@@ -54,8 +54,7 @@ pub enum QuerySideEffect {
 
 #[derive(Clone)]
 pub struct DepGraph {
-    is_in_sandbox: Arc<AtomicBool>,
-    data: [Option<Arc<DepGraphData>>; 2],
+    data: Arc<AtomicPtr<Option<DepGraphData>>>,
 
     /// This field is used for assigning DepNodeIndices when running in
     /// non-incremental mode. Even in non-incremental mode we make sure that
@@ -170,31 +169,26 @@ impl DepGraph {
         }
 
         DepGraph {
-            data: [
-                Some(Arc::new(DepGraphData {
-                    previous_work_products: prev_work_products,
-                    current,
-                    previous: prev_graph,
-                    colors,
-                    debug_loaded_from_disk: Default::default(),
-                })),
-                None,
-            ],
+            data: Arc::new(AtomicPtr::new(Box::into_raw(Box::new(Some(DepGraphData {
+                previous_work_products: prev_work_products,
+                current,
+                previous: prev_graph,
+                colors,
+                debug_loaded_from_disk: Default::default(),
+            }))))),
             virtual_dep_node_index: Arc::new(AtomicU32::new(0)),
-            is_in_sandbox: Default::default(),
         }
     }
 
     pub fn new_disabled() -> DepGraph {
         DepGraph {
-            data: [None, None],
+            data: Arc::new(AtomicPtr::new(Box::into_raw(Box::new(None)))),
             virtual_dep_node_index: Arc::new(AtomicU32::new(0)),
-            is_in_sandbox: Default::default(),
         }
     }
 
     pub fn data(&self) -> Option<&DepGraphData> {
-        self.data[self.is_in_sandbox.load(Ordering::Relaxed) as usize].as_deref()
+        unsafe { self.data.load(Ordering::Relaxed).as_ref_unchecked().as_ref() }
     }
 
     /// Returns `true` if we are actually building the full dep-graph, and `false` otherwise.
@@ -291,11 +285,12 @@ impl DepGraph {
 
     pub fn with_sandbox(&self, op: impl FnOnce()) {
         self.with_query_deserialization(|| {
-            self.is_in_sandbox.store(true, Ordering::Relaxed);
+            let real_graph = self.data.swap(Box::into_raw(Box::new(None)), Ordering::Relaxed);
 
             op();
 
-            self.is_in_sandbox.store(false, Ordering::Relaxed);
+            let fake_graph = self.data.swap(real_graph, Ordering::Relaxed);
+            unsafe { drop(Box::from_raw(fake_graph)) };
         });
     }
 

compiler/rustc_middle/src/query/caches.rs

@@ -1,5 +1,7 @@
+use std::sync::OnceLock;
+use std::sync::atomic::{AtomicPtr, Ordering};
+
 use rustc_data_structures::sharded::ShardedHashMap;
-use rustc_data_structures::sync::Lock;
 pub use rustc_data_structures::vec_cache::VecCache;
 use rustc_hir::def_id::LOCAL_CRATE;
 use rustc_index::Idx;
@@ -98,12 +100,12 @@ where
 /// In-memory cache for queries whose key type only has one value (e.g. `()`).
 /// The cache therefore only needs to store one query return value.
 pub struct SingleCache<V> {
-    cache: Lock<Option<(V, DepNodeIndex)>>,
+    cache: AtomicPtr<OnceLock<(V, DepNodeIndex)>>,
 }
 
 impl<V> Default for SingleCache<V> {
     fn default() -> Self {
-        SingleCache { cache: Default::default() }
+        SingleCache { cache: AtomicPtr::new(Box::into_raw(Box::new(Default::default()))) }
     }
 }
 
@@ -116,27 +118,33 @@ where
 
     #[inline(always)]
     fn lookup(&self, _key: &()) -> Option<(V, DepNodeIndex)> {
-        self.cache.lock().clone()
+        unsafe { self.cache.load(Ordering::Relaxed).as_ref_unchecked().get().copied() }
     }
 
     #[inline]
     fn complete(&self, _key: (), value: V, index: DepNodeIndex) {
-        *self.cache.lock() = Some((value, index))
+        unsafe {
+            self.cache.load(Ordering::Relaxed).as_ref_unchecked().set((value, index)).ok();
+        }
     }
 
     fn for_each(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
-        if let Some(value) = self.cache.lock().clone() {
-            f(&(), &value.0, value.1)
+        unsafe {
+            if let Some(value) = self.cache.load(Ordering::Relaxed).as_ref_unchecked().get() {
+                f(&(), &value.0, value.1)
+            }
         }
     }
 
     fn len(&self) -> usize {
-        self.cache.lock().is_some().into()
+        unsafe { self.cache.load(Ordering::Relaxed).as_ref_unchecked().get().is_some().into() }
     }
 
     fn invalidate(&self, selector: impl Fn(Self::Key) -> bool) {
         if selector(()) {
-            *self.cache.lock() = None;
+            let prev =
+                self.cache.swap(Box::into_raw(Box::new(Default::default())), Ordering::Relaxed);
+            unsafe { drop(Box::from_raw(prev)) }
         }
     }
 }