Store callbacks in global statics

The callbacks have precisely two states: the default, and the one
present throughout almost all of the rustc run (the filled in value
which has access to TyCtxt).

We used to store this as a thread local, and reset it on each thread to
the non-default value. But this is somewhat wasteful, since there is no
reason to set it globally -- while the callbacks themselves access TLS,
they do not do so in a manner that fails in when we do not have TLS to
work with.

Author: Mark-Simulacrum

src/librustc/ty/context.rs

@@ -1623,13 +1623,11 @@ pub mod tls {
 
     use crate::dep_graph::TaskDeps;
     use crate::ty::query;
-    use errors::{Diagnostic, TRACK_DIAGNOSTICS};
+    use errors::Diagnostic;
     use rustc_data_structures::sync::{self, Lock, Lrc};
     use rustc_data_structures::thin_vec::ThinVec;
     use rustc_data_structures::OnDrop;
-    use std::fmt;
     use std::mem;
-    use syntax_pos;
 
     #[cfg(not(parallel_compiler))]
     use std::cell::Cell;
@@ -1705,58 +1703,6 @@ pub mod tls {
         TLV.with(|tlv| tlv.get())
     }
 
-    /// This is a callback from libsyntax as it cannot access the implicit state
-    /// in librustc otherwise.
-    fn span_debug(span: syntax_pos::Span, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        with_opt(|tcx| {
-            if let Some(tcx) = tcx {
-                write!(f, "{}", tcx.sess.source_map().span_to_string(span))
-            } else {
-                syntax_pos::default_span_debug(span, f)
-            }
-        })
-    }
-
-    /// This is a callback from libsyntax as it cannot access the implicit state
-    /// in librustc otherwise. It is used to when diagnostic messages are
-    /// emitted and stores them in the current query, if there is one.
-    fn track_diagnostic(diagnostic: &Diagnostic) {
-        with_context_opt(|icx| {
-            if let Some(icx) = icx {
-                if let Some(ref diagnostics) = icx.diagnostics {
-                    let mut diagnostics = diagnostics.lock();
-                    diagnostics.extend(Some(diagnostic.clone()));
-                }
-            }
-        })
-    }
-
-    /// Sets up the callbacks from libsyntax on the current thread.
-    pub fn with_thread_locals<F, R>(f: F) -> R
-    where
-        F: FnOnce() -> R,
-    {
-        syntax_pos::SPAN_DEBUG.with(|span_dbg| {
-            let original_span_debug = span_dbg.get();
-            span_dbg.set(span_debug);
-
-            let _on_drop = OnDrop(move || {
-                span_dbg.set(original_span_debug);
-            });
-
-            TRACK_DIAGNOSTICS.with(|current| {
-                let original = current.get();
-                current.set(track_diagnostic);
-
-                let _on_drop = OnDrop(move || {
-                    current.set(original);
-                });
-
-                f()
-            })
-        })
-    }
-
     /// Sets `context` as the new current `ImplicitCtxt` for the duration of the function `f`.
     #[inline]
     pub fn enter_context<'a, 'tcx, F, R>(context: &ImplicitCtxt<'a, 'tcx>, f: F) -> R

src/librustc/ty/query/job.rs

@@ -438,9 +438,8 @@ pub unsafe fn handle_deadlock() {
     thread::spawn(move || {
         tls::GCX_PTR.set(gcx_ptr, || {
             syntax_pos::GLOBALS.set(syntax_pos_globals, || {
-                syntax_pos::GLOBALS.set(syntax_pos_globals, || {
-                    tls::with_thread_locals(|| tls::with_global(|tcx| deadlock(tcx, &registry)))
-                })
+                syntax_pos::GLOBALS
+                    .set(syntax_pos_globals, || tls::with_global(|tcx| deadlock(tcx, &registry)))
             })
         })
     });

src/librustc_data_structures/atomic_ref.rs

@@ -0,0 +1,26 @@
+use std::marker::PhantomData;
+use std::sync::atomic::{AtomicPtr, Ordering};
+
+/// This is essentially an `AtomicPtr` but is guaranteed to always be valid
+pub struct AtomicRef<T: 'static>(AtomicPtr<T>, PhantomData<&'static T>);
+
+impl<T: 'static> AtomicRef<T> {
+    pub const fn new(initial: &'static T) -> AtomicRef<T> {
+        AtomicRef(AtomicPtr::new(initial as *const T as *mut T), PhantomData)
+    }
+
+    pub fn swap(&self, new: &'static T) -> &'static T {
+        // We never allow storing anything but a `'static` reference so it's safe to
+        // return it for the same.
+        unsafe { &*self.0.swap(new as *const T as *mut T, Ordering::SeqCst) }
+    }
+}
+
+impl<T: 'static> std::ops::Deref for AtomicRef<T> {
+    type Target = T;
+    fn deref(&self) -> &Self::Target {
+        // We never allow storing anything but a `'static` reference so it's safe to lend
+        // it out for any amount of time.
+        unsafe { &*self.0.load(Ordering::SeqCst) }
+    }
+}

src/librustc_data_structures/lib.rs

@@ -89,10 +89,12 @@ pub mod thin_vec;
 pub mod tiny_list;
 pub mod transitive_relation;
 pub use ena::unify;
+mod atomic_ref;
 pub mod fingerprint;
 pub mod profiling;
 pub mod vec_linked_list;
 pub mod work_queue;
+pub use atomic_ref::AtomicRef;
 
 pub struct OnDrop<F: Fn()>(pub F);
 

src/librustc_errors/lib.rs

@@ -17,11 +17,11 @@ use registry::Registry;
 use rustc_data_structures::fx::{FxHashSet, FxIndexMap};
 use rustc_data_structures::stable_hasher::StableHasher;
 use rustc_data_structures::sync::{self, Lock, Lrc};
+use rustc_data_structures::AtomicRef;
 use syntax_pos::source_map::SourceMap;
 use syntax_pos::{Loc, MultiSpan, Span};
 
 use std::borrow::Cow;
-use std::cell::Cell;
 use std::panic;
 use std::path::Path;
 use std::{error, fmt};
@@ -313,8 +313,8 @@ pub enum StashKey {
 
 fn default_track_diagnostic(_: &Diagnostic) {}
 
-thread_local!(pub static TRACK_DIAGNOSTICS: Cell<fn(&Diagnostic)> =
-                Cell::new(default_track_diagnostic));
+pub static TRACK_DIAGNOSTICS: AtomicRef<fn(&Diagnostic)> =
+    AtomicRef::new(&(default_track_diagnostic as fn(&_)));
 
 #[derive(Copy, Clone, Default)]
 pub struct HandlerFlags {
@@ -734,9 +734,7 @@ impl HandlerInner {
             return;
         }
 
-        TRACK_DIAGNOSTICS.with(|track_diagnostics| {
-            track_diagnostics.get()(diagnostic);
-        });
+        (*TRACK_DIAGNOSTICS)(diagnostic);
 
         if let Some(ref code) = diagnostic.code {
             self.emitted_diagnostic_codes.insert(code.clone());

src/librustc_interface/callbacks.rs

@@ -0,0 +1,48 @@
+//! Throughout the compiler tree, there are several places which want to have
+//! access to state or queries while being inside crates that are dependencies
+//! of librustc. To facilitate this, we have the
+//! `rustc_data_structures::AtomicRef` type, which allows us to setup a global
+//! static which can then be set in this file at program startup.
+//!
+//! See `SPAN_DEBUG` for an example of how to set things up.
+//!
+//! The functions in this file should fall back to the default set in their
+//! origin crate when the `TyCtxt` is not present in TLS.
+
+use rustc::ty::tls;
+use rustc_errors::{Diagnostic, TRACK_DIAGNOSTICS};
+use std::fmt;
+use syntax_pos;
+
+/// This is a callback from libsyntax as it cannot access the implicit state
+/// in librustc otherwise.
+fn span_debug(span: syntax_pos::Span, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+    tls::with_opt(|tcx| {
+        if let Some(tcx) = tcx {
+            write!(f, "{}", tcx.sess.source_map().span_to_string(span))
+        } else {
+            syntax_pos::default_span_debug(span, f)
+        }
+    })
+}
+
+/// This is a callback from libsyntax as it cannot access the implicit state
+/// in librustc otherwise. It is used to when diagnostic messages are
+/// emitted and stores them in the current query, if there is one.
+fn track_diagnostic(diagnostic: &Diagnostic) {
+    tls::with_context_opt(|icx| {
+        if let Some(icx) = icx {
+            if let Some(ref diagnostics) = icx.diagnostics {
+                let mut diagnostics = diagnostics.lock();
+                diagnostics.extend(Some(diagnostic.clone()));
+            }
+        }
+    })
+}
+
+/// Sets up the callbacks in prior crates which we want to refer to the
+/// TyCtxt in.
+pub fn setup_callbacks() {
+    syntax_pos::SPAN_DEBUG.swap(&(span_debug as fn(_, &mut fmt::Formatter<'_>) -> _));
+    TRACK_DIAGNOSTICS.swap(&(track_diagnostic as fn(&_)));
+}

src/librustc_interface/lib.rs

@@ -11,6 +11,7 @@
 #[cfg(unix)]
 extern crate libc;
 
+mod callbacks;
 pub mod interface;
 mod passes;
 mod proc_macro_decls;

src/librustc_interface/util.rs

@@ -145,13 +145,15 @@ pub fn spawn_thread_pool<F: FnOnce() -> R + Send, R: Send>(
         cfg = cfg.stack_size(size);
     }
 
+    crate::callbacks::setup_callbacks();
+
     scoped_thread(cfg, || {
         syntax::with_globals(edition, || {
             ty::tls::GCX_PTR.set(&Lock::new(0), || {
                 if let Some(stderr) = stderr {
                     io::set_panic(Some(box Sink(stderr.clone())));
                 }
-                ty::tls::with_thread_locals(|| f())
+                f()
             })
         })
     })
@@ -167,6 +169,7 @@ pub fn spawn_thread_pool<F: FnOnce() -> R + Send, R: Send>(
     use rayon::{ThreadBuilder, ThreadPool, ThreadPoolBuilder};
 
     let gcx_ptr = &Lock::new(0);
+    crate::callbacks::setup_callbacks();
 
     let mut config = ThreadPoolBuilder::new()
         .thread_name(|_| "rustc".to_string())
@@ -194,9 +197,7 @@ pub fn spawn_thread_pool<F: FnOnce() -> R + Send, R: Send>(
                             if let Some(stderr) = stderr {
                                 io::set_panic(Some(box Sink(stderr.clone())));
                             }
-                            ty::tls::with_thread_locals(|| {
-                                ty::tls::GCX_PTR.set(gcx_ptr, || thread.run())
-                            })
+                            ty::tls::GCX_PTR.set(gcx_ptr, || thread.run())
                         })
                     })
                 };

src/libsyntax_pos/lib.rs

@@ -13,6 +13,7 @@
 #![feature(specialization)]
 #![feature(step_trait)]
 
+use rustc_data_structures::AtomicRef;
 use rustc_macros::HashStable_Generic;
 use rustc_serialize::{Decodable, Decoder, Encodable, Encoder};
 
@@ -41,7 +42,7 @@ use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync::{Lock, Lrc};
 
 use std::borrow::Cow;
-use std::cell::{Cell, RefCell};
+use std::cell::RefCell;
 use std::cmp::{self, Ordering};
 use std::fmt;
 use std::hash::{Hash, Hasher};
@@ -665,13 +666,13 @@ pub fn default_span_debug(span: Span, f: &mut fmt::Formatter<'_>) -> fmt::Result
 
 impl fmt::Debug for Span {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        SPAN_DEBUG.with(|span_debug| span_debug.get()(*self, f))
+        (*SPAN_DEBUG)(*self, f)
     }
 }
 
 impl fmt::Debug for SpanData {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        SPAN_DEBUG.with(|span_debug| span_debug.get()(Span::new(self.lo, self.hi, self.ctxt), f))
+        (*SPAN_DEBUG)(Span::new(self.lo, self.hi, self.ctxt), f)
     }
 }
 
@@ -1503,8 +1504,8 @@ pub struct FileLines {
     pub lines: Vec<LineInfo>,
 }
 
-thread_local!(pub static SPAN_DEBUG: Cell<fn(Span, &mut fmt::Formatter<'_>) -> fmt::Result> =
-                Cell::new(default_span_debug));
+pub static SPAN_DEBUG: AtomicRef<fn(Span, &mut fmt::Formatter<'_>) -> fmt::Result> =
+    AtomicRef::new(&(default_span_debug as fn(_, &mut fmt::Formatter<'_>) -> _));
 
 #[derive(Debug)]
 pub struct MacroBacktrace {