Use the extension mechanism for AnalysisCtxt

Author: nbdd0121

src/atomic_context.rs

@@ -298,15 +298,7 @@ impl<'tcx> AnalysisCtxt<'tcx> {
 }
 
 pub struct AtomicContext<'tcx> {
-    cx: AnalysisCtxt<'tcx>,
-}
-
-impl<'tcx> AtomicContext<'tcx> {
-    pub fn new(tcx: TyCtxt<'tcx>) -> Self {
-        Self {
-            cx: AnalysisCtxt::new(tcx),
-        }
-    }
+    pub cx: &'tcx AnalysisCtxt<'tcx>,
 }
 
 impl_lint_pass!(AtomicContext<'_> => [ATOMIC_CONTEXT]);
@@ -440,7 +432,7 @@ impl<'tcx> LateLintPass<'tcx> for AtomicContext<'tcx> {
             .tcx
             .erase_and_anonymize_regions(GenericArgs::identity_for_item(self.cx.tcx, def_id));
         let instance = Instance::new_raw(def_id.into(), identity);
-        let poly_instance = TypingEnv::post_analysis(*self.cx, def_id).as_query_input(instance);
+        let poly_instance = TypingEnv::post_analysis(self.cx.tcx, def_id).as_query_input(instance);
         let _ = self.cx.instance_adjustment(poly_instance);
         let _ = self.cx.instance_expectation(poly_instance);
         let _ = self.cx.instance_check(poly_instance);

src/ctxt.rs

@@ -3,11 +3,11 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 
 use std::any::{Any, TypeId};
-use std::cell::RefCell;
 use std::sync::Arc;
 
 use rusqlite::{Connection, OptionalExtension};
 use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::sync::{DynSend, DynSync, MTLock, RwLock};
 use rustc_hir::def_id::{CrateNum, LOCAL_CRATE};
 use rustc_middle::ty::TyCtxt;
 use rustc_serialize::{Decodable, Encodable};
@@ -18,8 +18,8 @@ use crate::preempt_count::UseSite;
 pub(crate) trait Query: 'static {
     const NAME: &'static str;
 
-    type Key<'tcx>;
-    type Value<'tcx>;
+    type Key<'tcx>: DynSend + DynSync;
+    type Value<'tcx>: DynSend + DynSync;
 }
 
 pub(crate) trait QueryValueDecodable: Query {
@@ -50,13 +50,17 @@ pub(crate) trait PersistentQuery: QueryValueDecodable {
 
 pub struct AnalysisCtxt<'tcx> {
     pub tcx: TyCtxt<'tcx>,
-    pub local_conn: Connection,
-    pub sql_conn: RefCell<FxHashMap<CrateNum, Option<Arc<Connection>>>>,
+    pub local_conn: MTLock<Connection>,
+    pub sql_conn: RwLock<FxHashMap<CrateNum, Option<Arc<MTLock<Connection>>>>>,
 
-    pub call_stack: RefCell<Vec<UseSite<'tcx>>>,
-    pub query_cache: RefCell<FxHashMap<TypeId, Arc<dyn Any>>>,
+    pub call_stack: RwLock<Vec<UseSite<'tcx>>>,
+    pub query_cache: RwLock<FxHashMap<TypeId, Arc<dyn Any + DynSend + DynSync>>>,
 }
 
+// Everything in `AnalysisCtxt` is either `DynSend/DynSync` or `Send/Sync`, but since there're no relation between two right now compiler cannot infer this.
+unsafe impl<'tcx> DynSend for AnalysisCtxt<'tcx> {}
+unsafe impl<'tcx> DynSync for AnalysisCtxt<'tcx> {}
+
 impl<'tcx> std::ops::Deref for AnalysisCtxt<'tcx> {
     type Target = TyCtxt<'tcx>;
 
@@ -105,7 +109,7 @@ const SCHEMA_VERSION: u32 = 1;
 
 impl Drop for AnalysisCtxt<'_> {
     fn drop(&mut self) {
-        self.local_conn.execute("commit", ()).unwrap();
+        self.local_conn.lock().execute("commit", ()).unwrap();
     }
 }
 
@@ -127,26 +131,26 @@ impl ArcDowncast for Arc<dyn Any> {
 impl<'tcx> AnalysisCtxt<'tcx> {
     pub(crate) fn query_cache<Q: Query>(
         &self,
-    ) -> Arc<RefCell<FxHashMap<Q::Key<'tcx>, Q::Value<'tcx>>>> {
+    ) -> Arc<RwLock<FxHashMap<Q::Key<'tcx>, Q::Value<'tcx>>>> {
         let key = TypeId::of::<Q>();
         let mut guard = self.query_cache.borrow_mut();
-        let cache = guard
+        let cache = (guard
             .entry(key)
             .or_insert_with(|| {
-                let cache = Arc::new(RefCell::new(
+                let cache = Arc::new(RwLock::new(
                     FxHashMap::<Q::Key<'static>, Q::Value<'static>>::default(),
                 ));
                 cache
             })
-            .clone()
-            .downcast::<RefCell<FxHashMap<Q::Key<'static>, Q::Value<'static>>>>()
+            .clone() as Arc<dyn Any>)
+            .downcast::<RwLock<FxHashMap<Q::Key<'static>, Q::Value<'static>>>>()
             .unwrap();
         // Everything stored inside query_cache is conceptually `'tcx`, but due to limitation
         // of `Any` we hack around the lifetime.
         unsafe { std::mem::transmute(cache) }
     }
 
-    pub(crate) fn sql_connection(&self, cnum: CrateNum) -> Option<Arc<Connection>> {
+    pub(crate) fn sql_connection(&self, cnum: CrateNum) -> Option<Arc<MTLock<Connection>>> {
         if let Some(v) = self.sql_conn.borrow().get(&cnum) {
             return v.clone();
         }
@@ -187,7 +191,7 @@ impl<'tcx> AnalysisCtxt<'tcx> {
                     );
                 }
 
-                result = Some(Arc::new(conn));
+                result = Some(Arc::new(MTLock::new(conn)));
                 break;
             }
         }
@@ -205,6 +209,7 @@ impl<'tcx> AnalysisCtxt<'tcx> {
 
     pub(crate) fn sql_create_table<Q: Query>(&self) {
         self.local_conn
+            .lock()
             .execute_batch(&format!(
                 "CREATE TABLE {} (key BLOB PRIMARY KEY, value BLOB);",
                 Q::NAME
@@ -225,6 +230,7 @@ impl<'tcx> AnalysisCtxt<'tcx> {
 
         let value_encoded: Vec<u8> = self
             .sql_connection(cnum)?
+            .lock()
             .query_row(
                 &format!("SELECT value FROM {} WHERE key = ?", Q::NAME),
                 rusqlite::params![encoded],
@@ -265,6 +271,7 @@ impl<'tcx> AnalysisCtxt<'tcx> {
         let value_encoded = encode_ctx.finish();
 
         self.local_conn
+            .lock()
             .execute(
                 &format!(
                     "INSERT OR REPLACE INTO {} (key, value) VALUES (?, ?)",
@@ -309,7 +316,7 @@ impl<'tcx> AnalysisCtxt<'tcx> {
 
         let ret = Self {
             tcx,
-            local_conn: conn,
+            local_conn: MTLock::new(conn),
             sql_conn: Default::default(),
             call_stack: Default::default(),
             query_cache: Default::default(),

src/main.rs

@@ -48,6 +48,8 @@ use rustc_session::EarlyDiagCtxt;
 use rustc_session::config::ErrorOutputType;
 use std::sync::atomic::Ordering;
 
+use crate::ctxt::AnalysisCtxt;
+
 #[macro_use]
 mod ctxt;
 
@@ -78,14 +80,16 @@ impl Callbacks for MyCallbacks {
         config.override_queries = Some(|_, provider| {
             // Calling `optimized_mir` will steal the result of query `mir_drops_elaborated_and_const_checked`,
             // so hijack `optimized_mir` to run `analysis_mir` first.
-            hook_query!(provider.optimized_mir => |tcx, def_id, original| {
+            hook_query!(provider.optimized_mir => |tcx, local_def_id, original| {
+                let def_id = local_def_id.to_def_id();
                 // Skip `analysis_mir` call if this is a constructor, since it will be delegated back to
                 // `optimized_mir` for building ADT constructor shim.
-                if !tcx.is_constructor(def_id.to_def_id()) {
-                    crate::mir::local_analysis_mir(tcx, def_id);
+                if !tcx.is_constructor(def_id) {
+                    let cx = crate::driver::cx::<MyCallbacks>(tcx);
+                    let _ = cx.analysis_mir(def_id);
                 }
 
-                original(tcx, def_id)
+                original(tcx, local_def_id)
             });
         });
         config.register_lints = Some(Box::new(move |_, lint_store| {
@@ -94,16 +98,20 @@ impl Callbacks for MyCallbacks {
             lint_store.register_lints(&[&atomic_context::ATOMIC_CONTEXT]);
             // lint_store
             //     .register_late_pass(|_| Box::new(infallible_allocation::InfallibleAllocation));
-            lint_store.register_late_pass(|tcx| Box::new(atomic_context::AtomicContext::new(tcx)));
+            lint_store.register_late_pass(|tcx| {
+                Box::new(atomic_context::AtomicContext {
+                    cx: driver::cx::<MyCallbacks>(tcx),
+                })
+            });
         }));
     }
 }
 
 impl driver::CallbacksExt for MyCallbacks {
-    type ExtCtxt<'tcx> = TyCtxt<'tcx>;
+    type ExtCtxt<'tcx> = AnalysisCtxt<'tcx>;
 
     fn ext_cx<'tcx>(&mut self, tcx: TyCtxt<'tcx>) -> Self::ExtCtxt<'tcx> {
-        tcx
+        AnalysisCtxt::new(tcx)
     }
 }
 

src/mir.rs

@@ -6,11 +6,6 @@ pub mod drop_shim;
 pub mod elaborate_drop;
 pub mod patch;
 
-use std::sync::atomic::AtomicPtr;
-use std::sync::atomic::Ordering;
-use std::sync::{LazyLock, Mutex};
-
-use rustc_data_structures::fx::FxHashMap;
 use rustc_hir::{self as hir, def::DefKind};
 use rustc_middle::mir::CallSource;
 use rustc_middle::mir::{
@@ -24,38 +19,17 @@ use rustc_span::{DUMMY_SP, source_map::Spanned, sym};
 use crate::ctxt::AnalysisCtxt;
 use crate::ctxt::PersistentQuery;
 
-// HACK: we can't add new queries to `TyCtxt` without changing rustc code, so
-// use this as a "poor man's query" for now.
-//
-// `AnalysisCtxt` has its own cache but we can't use it as we've only got `TyCtxt`
-// so far but not `AnalysisCtxt`.
-static MIR_CACHE: LazyLock<Mutex<FxHashMap<LocalDefId, AtomicPtr<()>>>> =
-    LazyLock::new(|| Mutex::new(FxHashMap::default()));
-
-pub fn local_analysis_mir<'tcx>(tcx: TyCtxt<'tcx>, did: LocalDefId) -> &'tcx Body<'tcx> {
-    if tcx.is_constructor(did.to_def_id()) {
-        return tcx.optimized_mir(did.to_def_id());
+pub fn local_analysis_mir<'tcx>(cx: &AnalysisCtxt<'tcx>, did: LocalDefId) -> &'tcx Body<'tcx> {
+    if cx.is_constructor(did.to_def_id()) {
+        return cx.optimized_mir(did.to_def_id());
     }
 
-    {
-        let cache = MIR_CACHE.lock().unwrap();
-        if let Some(body) = cache.get(&did) {
-            return unsafe { &*body.load(Ordering::Relaxed).cast() };
-        }
-    }
-
-    let body = tcx
+    let body = cx
         .mir_drops_elaborated_and_const_checked(did)
         .borrow()
         .clone();
-    let body = remap_mir_for_const_eval_select(tcx, body, hir::Constness::NotConst);
-    let body = tcx.arena.alloc(body);
-
-    {
-        let mut cache = MIR_CACHE.lock().unwrap();
-        cache.insert(did, AtomicPtr::new(body as *const _ as *mut _));
-    }
-    body
+    let body = remap_mir_for_const_eval_select(cx.tcx, body, hir::Constness::NotConst);
+    cx.arena.alloc(body)
 }
 
 // Copied from rustc_mir_transform/src/lib.rs.
@@ -148,7 +122,7 @@ fn take_array<T, const N: usize>(b: &mut Box<[T]>) -> Result<[T; N], Box<[T]>> {
 memoize!(
     pub fn analysis_mir<'tcx>(cx: &AnalysisCtxt<'tcx>, def_id: DefId) -> &'tcx Body<'tcx> {
         if let Some(local_def_id) = def_id.as_local() {
-            local_analysis_mir(cx.tcx, local_def_id)
+            local_analysis_mir(cx, local_def_id)
         } else if let Some(mir) = cx.sql_load_with_span::<analysis_mir>(def_id, cx.def_span(def_id))
         {
             mir