Auto merge of rust-lang#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.

Author: bors

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,
             }
         }

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, || {

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;

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())
                         ()
                     ),)*
                 }