Refactor anymap to its own module

Author: nbdd0121

.vscode/settings.json

@@ -1,3 +1,4 @@
 {
-    "rust-analyzer.rustc.source": "discover"
+    "rust-analyzer.rustc.source": "discover",
+    "rust-analyzer.imports.granularity.group": "module"
 }

src/ctxt.rs

@@ -2,7 +2,8 @@
 //
 // SPDX-License-Identifier: MIT OR Apache-2.0
 
-use std::any::{Any, TypeId};
+use std::any::Any;
+use std::marker::PhantomData;
 use std::sync::Arc;
 
 use rusqlite::{Connection, OptionalExtension};
@@ -15,6 +16,7 @@ use rustc_session::config::OutputType;
 use rustc_span::{DUMMY_SP, Span};
 
 use crate::diagnostic::use_stack::UseSite;
+use crate::utils::anymap::AnyMap;
 
 pub(crate) trait Query: 'static {
     const NAME: &'static str;
@@ -55,7 +57,7 @@ pub struct AnalysisCtxt<'tcx> {
     pub sql_conn: RwLock<FxHashMap<CrateNum, Option<Arc<MTLock<Connection>>>>>,
 
     pub call_stack: RwLock<Vec<UseSite<'tcx>>>,
-    pub query_cache: RwLock<FxHashMap<TypeId, Arc<dyn Any + DynSend + DynSync>>>,
+    pub query_cache: RwLock<AnyMap<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.
@@ -114,38 +116,21 @@ impl Drop for AnalysisCtxt<'_> {
     }
 }
 
-// Used when parallel compiler is used.
-trait ArcDowncast: Sized {
-    fn downcast<T: Any>(self) -> Result<Arc<T>, Self>;
-}
-
-impl ArcDowncast for Arc<dyn Any> {
-    fn downcast<T: Any>(self) -> Result<Arc<T>, Self> {
-        if (*self).is::<T>() {
-            Ok(unsafe { Arc::from_raw(Arc::into_raw(self) as _) })
-        } else {
-            Err(self)
-        }
-    }
-}
-
 impl<'tcx> AnalysisCtxt<'tcx> {
     pub(crate) fn query_cache<Q: Query>(
         &self,
     ) -> 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
-            .entry(key)
+        let cache = guard
+            .entry()
             .or_insert_with(|| {
                 let cache = Arc::new(RwLock::new(
                     FxHashMap::<Q::Key<'static>, Q::Value<'static>>::default(),
                 ));
-                cache
+                (PhantomData::<fn() -> Q>, cache)
             })
-            .clone() as Arc<dyn Any>)
-            .downcast::<RwLock<FxHashMap<Q::Key<'static>, Q::Value<'static>>>>()
-            .unwrap();
+            .1
+            .clone();
         // Everything stored inside query_cache is conceptually `'tcx`, but due to limitation
         // of `Any` we hack around the lifetime.
         unsafe { std::mem::transmute(cache) }

src/main.rs

@@ -12,6 +12,7 @@
 #![feature(once_cell_get_mut)]
 // Used in symbol.rs
 #![feature(macro_metavar_expr)]
+#![feature(unsize)]
 #![warn(rustc::internal)]
 
 #[macro_use]
@@ -70,6 +71,7 @@ mod preempt_count;
 mod serde;
 mod symbol;
 mod util;
+mod utils;
 
 rustc_session::declare_tool_lint! {
     pub klint::INCORRECT_ATTRIBUTE,

src/utils/anymap.rs

@@ -0,0 +1,92 @@
+use std::any::{Any, TypeId};
+use std::collections::hash_map as map;
+use std::marker::{PhantomData, Unsize};
+
+use rustc_data_structures::fx::FxHashMap;
+
+/// Map that can store data for arbitrary types.
+pub struct AnyMap<U: ?Sized> {
+    // This is basically `FxHashMap<TypeId, Box<dyn Any>>`
+    //
+    // The generic `U` is present to capture auto trait bounds.
+    map: FxHashMap<TypeId, Box<U>>,
+}
+
+pub struct OccupiedEntry<'a, U: ?Sized, T> {
+    entry: map::OccupiedEntry<'a, TypeId, Box<U>>,
+    phantom: PhantomData<T>,
+}
+
+impl<'a, U: Any + ?Sized + 'static, T: 'static> OccupiedEntry<'a, U, T> {
+    pub fn into_mut(self) -> &'a mut T
+    where
+        T: Unsize<U>,
+    {
+        let any_ref = &mut **self.entry.into_mut();
+        debug_assert_eq!((*any_ref).type_id(), TypeId::of::<T>());
+        // SAFETY: by type invariant, `any_ref` is a `&mut T`.
+        unsafe { &mut *(any_ref as *mut U as *mut T) }
+    }
+}
+
+pub struct VacantEntry<'a, U: ?Sized, T> {
+    entry: map::VacantEntry<'a, TypeId, Box<U>>,
+    phantom: PhantomData<T>,
+}
+
+impl<'a, U: Any + ?Sized, T> VacantEntry<'a, U, T> {
+    pub fn insert(self, value: T) -> &'a mut T
+    where
+        T: Unsize<U>,
+    {
+        let any_ref = &mut **self.entry.insert(Box::new(value) as _);
+        // SAFETY: we just inserted it and we know the type is `Box<T>`.
+        unsafe { &mut *(any_ref as *mut U as *mut T) }
+    }
+}
+
+pub enum Entry<'a, U: ?Sized, T> {
+    Occupied(OccupiedEntry<'a, U, T>),
+    Vacant(VacantEntry<'a, U, T>),
+}
+
+impl<'a, U: Any + ?Sized, T: 'static> Entry<'a, U, T> {
+    pub fn or_insert_with<F: FnOnce() -> T>(self, default: F) -> &'a mut T
+    where
+        T: Unsize<U>,
+    {
+        match self {
+            Entry::Occupied(entry) => entry.into_mut(),
+            Entry::Vacant(entry) => entry.insert(default()),
+        }
+    }
+}
+
+impl<U: ?Sized> Default for AnyMap<U> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<U: ?Sized> AnyMap<U> {
+    pub fn new() -> Self {
+        Self {
+            map: Default::default(),
+        }
+    }
+}
+
+impl<U: Any + ?Sized> AnyMap<U> {
+    pub fn entry<T: 'static>(&mut self) -> Entry<'_, U, T> {
+        match self.map.entry(TypeId::of::<T>()) {
+            map::Entry::Occupied(entry) => Entry::Occupied(OccupiedEntry {
+                entry,
+                phantom: PhantomData,
+            }),
+            map::Entry::Vacant(entry) => Entry::Vacant(VacantEntry {
+                entry,
+                phantom: PhantomData,
+            }),
+        }
+    }
+}

src/utils/mod.rs

@@ -0,0 +1 @@
+pub mod anymap;