Auto merge of #118227 - Mark-Simulacrum:worker-local-outline, r=cjgillot

Optimize QueryArena allocation

This shifts the WorkerLocal wrapper to be outside the QueryArena, meaning that instead of having each query allocate distinct arenas per-worker we allocate the full set of arenas per-worker. This is primarily a code size optimization (locally, ~85 kilobytes, [perf is reporting >100 kilobytes](https://perf.rust-lang.org/compare.html?start=1fd418f92ed13db88a21865ba5d909abcf16b6cc&end=884c95a3f1fe8d28630ec3cdb0c8f95b2e539fde&stat=instructions%3Au&tab=artifact-size)), saving a bunch of code in the initialization of the arenas which was previously duplicated lots of times (per arena type).

Additionally this tells LLVM that the thread count can't be zero in this code (I believe this is true?) which shaves some small amount of bytes off as well since we eliminate checks for zero in the vec allocations.

4 files changed, 12 insertions, 9 deletions

diff --git a/compiler/rustc_data_structures/src/sync/worker_local.rs b/compiler/rustc_data_structures/src/sync/worker_local.rs
index ffafdba13ce..b34d3dd9044 100644
--- a/compiler/rustc_data_structures/src/sync/worker_local.rs
+++ b/compiler/rustc_data_structures/src/sync/worker_local.rs
@@ -1,6 +1,7 @@
 use parking_lot::Mutex;
 use std::cell::Cell;
 use std::cell::OnceCell;
+use std::num::NonZeroUsize;
 use std::ops::Deref;
 use std::ptr;
 use std::sync::Arc;
@@ -30,7 +31,7 @@ impl RegistryId {
 }
 
 struct RegistryData {
-    thread_limit: usize,
+    thread_limit: NonZeroUsize,
     threads: Mutex<usize>,
 }
 
@@ -60,7 +61,7 @@ thread_local! {
 
 impl Registry {
     /// Creates a registry which can hold up to `thread_limit` threads.
-    pub fn new(thread_limit: usize) -> Self {
+    pub fn new(thread_limit: NonZeroUsize) -> Self {
         Registry(Arc::new(RegistryData { thread_limit, threads: Mutex::new(0) }))
     }
 
@@ -73,7 +74,7 @@ impl Registry {
     /// Panics if the thread limit is hit or if the thread already has an associated registry.
     pub fn register(&self) {
         let mut threads = self.0.threads.lock();
-        if *threads < self.0.thread_limit {
+        if *threads < self.0.thread_limit.get() {
             REGISTRY.with(|registry| {
                 if registry.get().is_some() {
                     drop(threads);
@@ -126,7 +127,9 @@ impl<T> WorkerLocal<T> {
         {
             let registry = Registry::current();
             WorkerLocal {
-                locals: (0..registry.0.thread_limit).map(|i| CacheAligned(initial(i))).collect(),
+                locals: (0..registry.0.thread_limit.get())
+                    .map(|i| CacheAligned(initial(i)))
+                    .collect(),
                 registry,
             }
         }
diff --git a/compiler/rustc_interface/src/util.rs b/compiler/rustc_interface/src/util.rs
index 8df3648a065..b3ab01a740a 100644
--- a/compiler/rustc_interface/src/util.rs
+++ b/compiler/rustc_interface/src/util.rs
@@ -107,7 +107,7 @@ pub(crate) fn run_in_thread_pool_with_globals<F: FnOnce() -> R + Send, R: Send>(
     use rustc_query_impl::QueryCtxt;
     use rustc_query_system::query::{deadlock, QueryContext};
 
-    let registry = sync::Registry::new(threads);
+    let registry = sync::Registry::new(std::num::NonZeroUsize::new(threads).unwrap());
 
     if !sync::is_dyn_thread_safe() {
         return run_in_thread_with_globals(edition, || {
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index c982e2a9325..22f3d733dc3 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -61,7 +61,6 @@ use rustc_data_structures::fx::{FxHashMap, FxIndexMap, FxIndexSet};
 use rustc_data_structures::steal::Steal;
 use rustc_data_structures::svh::Svh;
 use rustc_data_structures::sync::Lrc;
-use rustc_data_structures::sync::WorkerLocal;
 use rustc_data_structures::unord::UnordSet;
 use rustc_errors::ErrorGuaranteed;
 use rustc_hir as hir;
diff --git a/compiler/rustc_middle/src/query/plumbing.rs b/compiler/rustc_middle/src/query/plumbing.rs
index f4a8ada8f68..ac3538f181e 100644
--- a/compiler/rustc_middle/src/query/plumbing.rs
+++ b/compiler/rustc_middle/src/query/plumbing.rs
@@ -11,6 +11,7 @@ use field_offset::FieldOffset;
 use measureme::StringId;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::sync::AtomicU64;
+use rustc_data_structures::sync::WorkerLocal;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_hir::hir_id::OwnerId;
@@ -71,7 +72,7 @@ pub struct QuerySystemFns<'tcx> {
 
 pub struct QuerySystem<'tcx> {
     pub states: QueryStates<'tcx>,
-    pub arenas: QueryArenas<'tcx>,
+    pub arenas: WorkerLocal<QueryArenas<'tcx>>,
     pub caches: QueryCaches<'tcx>,
     pub dynamic_queries: DynamicQueries<'tcx>,
 
@@ -370,7 +371,7 @@ macro_rules! define_callbacks {
 
         pub struct QueryArenas<'tcx> {
             $($(#[$attr])* pub $name: query_if_arena!([$($modifiers)*]
-                (WorkerLocal<TypedArena<<$V as Deref>::Target>>)
+                (TypedArena<<$V as Deref>::Target>)
                 ()
             ),)*
         }
@@ -379,7 +380,7 @@ macro_rules! define_callbacks {
             fn default() -> Self {
                 Self {
                     $($name: query_if_arena!([$($modifiers)*]
-                        (WorkerLocal::new(|_| Default::default()))
+                        (Default::default())
                         ()
                     ),)*
                 }