Auto merge of #142277 - Mark-Simulacrum:active-query-by-idx, r=<try>

Key active queries by usize rather than storing raw keys

Starting with perf run to evaluate impact -- my hope is that this helps somewhat with rustc compile times, since it should allow de-duplicating a bunch of `HashTable<K, ...>` codegen since it's now ~always `HashTable<usize, ...>` -- more can be done by merging the DefaultCache's active + non-active maps too.

Locally I don't really see any win from this though, so it's possible it's wasted effort -- want to see what perf says about it.

r? ghost

Author: rust-bors[bot]

compiler/rustc_data_structures/src/sharded.rs

@@ -199,6 +199,17 @@ impl<K: Eq + Hash, V> ShardedHashMap<K, V> {
             }
         }
     }
+
+    #[inline]
+    pub fn remove(&self, key: K) -> Option<V> {
+        let hash = make_hash(&key);
+        let mut shard = self.lock_shard_by_hash(hash);
+
+        match table_entry(&mut shard, hash, &key) {
+            Entry::Occupied(e) => Some(e.remove().0.1),
+            Entry::Vacant(_) => None,
+        }
+    }
 }
 
 impl<K: Eq + Hash + Copy> ShardedHashMap<K, ()> {

compiler/rustc_data_structures/src/vec_cache.rs

@@ -8,6 +8,7 @@
 
 use std::fmt::Debug;
 use std::marker::PhantomData;
+use std::ptr::NonNull;
 use std::sync::atomic::{AtomicPtr, AtomicU32, AtomicUsize, Ordering};
 
 use rustc_index::Idx;
@@ -56,6 +57,8 @@ const ENTRIES_BY_BUCKET: [usize; 21] = {
     entries
 };
 
+const FIRST_BUCKET_SHIFT: usize = 12;
+
 impl SlotIndex {
     // This unpacks a flat u32 index into identifying which bucket it belongs to and the offset
     // within that bucket. As noted in the VecCache docs, buckets double in size with each index.
@@ -68,7 +71,6 @@ impl SlotIndex {
     // slots (see `slot_index_exhaustive` in tests).
     #[inline]
     const fn from_index(idx: u32) -> Self {
-        const FIRST_BUCKET_SHIFT: usize = 12;
         if idx < (1 << FIRST_BUCKET_SHIFT) {
             return SlotIndex {
                 bucket_idx: 0,
@@ -86,6 +88,15 @@ impl SlotIndex {
         }
     }
 
+    fn to_key<K: Idx>(&self) -> K {
+        let idx = if self.bucket_idx == 0 {
+            self.index_in_bucket
+        } else {
+            (1 << (self.bucket_idx + FIRST_BUCKET_SHIFT - 1)) + self.index_in_bucket
+        };
+        K::new(idx)
+    }
+
     // SAFETY: Buckets must be managed solely by functions here (i.e., get/put on SlotIndex) and
     // `self` comes from SlotIndex::from_index
     #[inline]
@@ -124,6 +135,22 @@ impl SlotIndex {
         Some((value, index))
     }
 
+    // SAFETY: Buckets must be managed solely by functions here (i.e., get/put on SlotIndex) and
+    // `self` comes from SlotIndex::from_index
+    #[inline]
+    unsafe fn addr<V>(&self, buckets: &[AtomicPtr<Slot<V>>; 21]) -> usize {
+        // SAFETY: `bucket_idx` is ilog2(u32).saturating_sub(11), which is at most 21, i.e.,
+        // in-bounds of buckets. See `from_index` for computation.
+        let bucket = unsafe { buckets.get_unchecked(self.bucket_idx) };
+        let ptr = self.bucket_ptr(bucket);
+        assert!(self.index_in_bucket < self.entries);
+        // SAFETY: `bucket` was allocated (so <= isize in total bytes) to hold `entries`, so this
+        // must be inbounds.
+        let slot = unsafe { ptr.add(self.index_in_bucket) };
+
+        slot.addr()
+    }
+
     fn bucket_ptr<V>(&self, bucket: &AtomicPtr<Slot<V>>) -> *mut Slot<V> {
         let ptr = bucket.load(Ordering::Acquire);
         if ptr.is_null() { self.initialize_bucket(bucket) } else { ptr }
@@ -331,6 +358,37 @@ where
             }
         }
     }
+
+    pub fn to_slot_address(&self, key: &K) -> usize {
+        let key = u32::try_from(key.index()).unwrap();
+        let slot_idx = SlotIndex::from_index(key);
+        // SAFETY: Same safety as lookup.
+        unsafe { slot_idx.addr(&self.buckets) }
+    }
+
+    pub fn to_key(&self, slot_addr: usize) -> K {
+        for (bucket_idx, bucket) in self.buckets.iter().enumerate() {
+            let Some(bucket_ptr) = NonNull::new(bucket.load(Ordering::Acquire)) else {
+                continue;
+            };
+            let entries = ENTRIES_BY_BUCKET[bucket_idx];
+
+            // Are we in this bucket?
+            if !(bucket_ptr.as_ptr().addr() <= slot_addr
+                && slot_addr <= bucket_ptr.as_ptr().wrapping_add(entries).addr())
+            {
+                continue;
+            }
+
+            let index_in_bucket =
+                (slot_addr - bucket_ptr.as_ptr().addr()) / std::mem::size_of::<Slot<V>>();
+            let slot_idx = SlotIndex { bucket_idx, entries, index_in_bucket };
+
+            return slot_idx.to_key::<K>();
+        }
+
+        unreachable!()
+    }
 }
 
 #[cfg(test)]

compiler/rustc_middle/src/query/plumbing.rs

@@ -491,7 +491,7 @@ macro_rules! define_callbacks {
         #[derive(Default)]
         pub struct QueryStates<'tcx> {
             $(
-                pub $name: QueryState<$($K)*, QueryStackDeferred<'tcx>>,
+                pub $name: QueryState<QueryStackDeferred<'tcx>>,
             )*
         }
 

compiler/rustc_query_impl/src/lib.rs

@@ -83,10 +83,7 @@ where
     }
 
     #[inline(always)]
-    fn query_state<'a>(
-        self,
-        qcx: QueryCtxt<'tcx>,
-    ) -> &'a QueryState<Self::Key, QueryStackDeferred<'tcx>>
+    fn query_state<'a>(self, qcx: QueryCtxt<'tcx>) -> &'a QueryState<QueryStackDeferred<'tcx>>
     where
         QueryCtxt<'tcx>: 'a,
     {
@@ -95,7 +92,7 @@ where
         unsafe {
             &*(&qcx.tcx.query_system.states as *const QueryStates<'tcx>)
                 .byte_add(self.dynamic.query_state)
-                .cast::<QueryState<Self::Key, QueryStackDeferred<'tcx>>>()
+                .cast::<QueryState<QueryStackDeferred<'tcx>>>()
         }
     }
 

compiler/rustc_query_impl/src/plumbing.rs

@@ -735,9 +735,10 @@ macro_rules! define_queries {
                 tcx: TyCtxt<'tcx>,
                 qmap: &mut QueryMap<QueryStackDeferred<'tcx>>,
             ) -> Option<()> {
-                let make_query = |tcx, key| {
+                let make_query = |tcx: TyCtxt<'tcx>, key| {
                     let kind = rustc_middle::dep_graph::dep_kinds::$name;
                     let name = stringify!($name);
+                    let key = tcx.query_system.caches.$name.to_key(key);
                     $crate::plumbing::create_query_frame(tcx, rustc_middle::query::descs::$name, key, kind, name)
                 };
                 let res = tcx.query_system.states.$name.try_collect_active_jobs(

compiler/rustc_query_system/src/query/caches.rs

@@ -1,6 +1,6 @@
 use std::fmt::Debug;
 use std::hash::Hash;
-use std::sync::OnceLock;
+use std::sync::{Mutex, OnceLock};
 
 use rustc_data_structures::sharded::ShardedHashMap;
 pub use rustc_data_structures::vec_cache::VecCache;
@@ -23,6 +23,17 @@ pub trait QueryCache: Sized {
     /// given key, if it is present in the cache.
     fn lookup(&self, key: &Self::Key) -> Option<(Self::Value, DepNodeIndex)>;
 
+    /// Start the job.
+    ///
+    /// After the first call to `to_unique_index`, the returned index is guaranteed to be
+    /// uniquely bound to `key` until earlier of `drop(self)` or `complete(key, ...)`.
+    fn to_unique_index(&self, key: &Self::Key) -> usize;
+
+    /// Reverse the mapping of to_unique_index to a key.
+    ///
+    /// Will panic if called with an index that's not currently available.
+    fn to_key(&self, idx: usize) -> Self::Key;
+
     /// Adds a key/value entry to this cache.
     ///
     /// Called by some part of the query system, after having obtained the
@@ -36,11 +47,12 @@ pub trait QueryCache: Sized {
 /// more specialized kinds of cache. Backed by a sharded hashmap.
 pub struct DefaultCache<K, V> {
     cache: ShardedHashMap<K, (V, DepNodeIndex)>,
+    active: Mutex<Vec<Option<K>>>,
 }
 
 impl<K, V> Default for DefaultCache<K, V> {
     fn default() -> Self {
-        DefaultCache { cache: Default::default() }
+        DefaultCache { cache: Default::default(), active: Default::default() }
     }
 }
 
@@ -57,11 +69,49 @@ where
         self.cache.get(key)
     }
 
+    fn to_unique_index(&self, key: &Self::Key) -> usize {
+        let mut guard = self.active.lock().unwrap();
+
+        for (idx, slot) in guard.iter_mut().enumerate() {
+            if let Some(k) = slot
+                && k == key
+            {
+                // Return idx if we found the slot containing this key.
+                return idx;
+            } else if slot.is_none() {
+                // If slot is empty, reuse it.
+                *slot = Some(*key);
+                return idx;
+            }
+        }
+
+        // If no slot currently contains our key, add a new slot.
+        let idx = guard.len();
+        guard.push(Some(*key));
+        return idx;
+    }
+
+    fn to_key(&self, idx: usize) -> Self::Key {
+        let guard = self.active.lock().unwrap();
+        guard[idx].expect("still present")
+    }
+
     #[inline]
     fn complete(&self, key: K, value: V, index: DepNodeIndex) {
         // We may be overwriting another value. This is all right, since the dep-graph
         // will check that the fingerprint matches.
         self.cache.insert(key, (value, index));
+
+        // Make sure to do this second -- this ensures lookups return success prior to active
+        // getting removed, helping avoiding assignment of multiple indices per logical key.
+        let mut guard = self.active.lock().unwrap();
+        for slot in guard.iter_mut() {
+            if let Some(k) = slot
+                && *k == key
+            {
+                *slot = None;
+            }
+        }
     }
 
     fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
@@ -102,6 +152,15 @@ where
         self.cache.set((value, index)).ok();
     }
 
+    fn to_unique_index(&self, _: &Self::Key) -> usize {
+        // SingleCache has a single key, so we can map directly to a constant.
+        0
+    }
+
+    fn to_key(&self, _: usize) -> Self::Key {
+        ()
+    }
+
     fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
         if let Some(value) = self.cache.get() {
             f(&(), &value.0, value.1)
@@ -142,6 +201,32 @@ where
         }
     }
 
+    fn to_unique_index(&self, key: &Self::Key) -> usize {
+        if key.krate == LOCAL_CRATE {
+            // The local cache assigns keys based on allocated addresses in the backing VecCache.
+            //
+            // Those addresses are always at least 4-aligned (due to DepNodeIndex), so the low bit
+            // can't be set. This means these don't overlap with the other cache given we |1 those
+            // IDs. We check this with the assertion.
+            let local_idx = self.local.to_unique_index(&key.index);
+            assert!(local_idx & 1 == 0);
+            local_idx
+        } else {
+            // Shifting is safe because DefaultCache uses only u32 for its indices, so we won't
+            // overflow here.
+            (self.foreign.to_unique_index(key) << 1) | 1
+        }
+    }
+
+    fn to_key(&self, idx: usize) -> Self::Key {
+        if idx & 1 == 0 {
+            // local
+            DefId::local(self.local.to_key(idx))
+        } else {
+            self.foreign.to_key(idx >> 1)
+        }
+    }
+
     #[inline]
     fn complete(&self, key: DefId, value: V, index: DepNodeIndex) {
         if key.krate == LOCAL_CRATE {
@@ -180,4 +265,12 @@ where
     fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
         self.iter(f)
     }
+
+    fn to_unique_index(&self, key: &Self::Key) -> usize {
+        self.to_slot_address(key)
+    }
+
+    fn to_key(&self, idx: usize) -> Self::Key {
+        self.to_key(idx)
+    }
 }

compiler/rustc_query_system/src/query/config.rs

@@ -28,7 +28,7 @@ pub trait QueryConfig<Qcx: QueryContext>: Copy {
     fn format_value(self) -> fn(&Self::Value) -> String;
 
     // Don't use this method to access query results, instead use the methods on TyCtxt
-    fn query_state<'a>(self, tcx: Qcx) -> &'a QueryState<Self::Key, Qcx::QueryInfo>
+    fn query_state<'a>(self, tcx: Qcx) -> &'a QueryState<Qcx::QueryInfo>
     where
         Qcx: 'a;