Auto merge of #49045 - Zoxc:tls, r=michaelwoerister

Make queries thread safe

This makes queries thread safe by removing the query stack and making queries point to their parents. Queries write to the query map when starting and cycles are detected by checking if there's already an entry in the query map. This makes cycle detection O(1) instead of O(n), where `n` is the size of the query stack.

This is mostly corresponds to the method I described [here](https://internals.rust-lang.org/t/parallelizing-rustc-using-rayon/6606).

cc @rust-lang/compiler

r? @michaelwoerister

src/librustc/ty/context.rs

@@ -1244,7 +1244,7 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
                                     Lrc::new(StableVec::new(v)));
         }
 
-        tls::enter_global(GlobalCtxt {
+        let gcx = &GlobalCtxt {
             sess: s,
             cstore,
             global_arenas: &arenas.global,
@@ -1285,7 +1285,9 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
             all_traits: RefCell::new(None),
             tx_to_llvm_workers: tx,
             output_filenames: Arc::new(output_filenames.clone()),
-       }, f)
+        };
+
+        tls::enter_global(gcx, f)
     }
 
     pub fn consider_optimizing<T: Fn() -> String>(&self, msg: T) -> bool {
@@ -1509,11 +1511,28 @@ impl<'a, 'tcx> TyCtxt<'a, 'tcx, 'tcx> {
 
 impl<'gcx: 'tcx, 'tcx> GlobalCtxt<'gcx> {
     /// Call the closure with a local `TyCtxt` using the given arena.
-    pub fn enter_local<F, R>(&self, arena: &'tcx DroplessArena, f: F) -> R
-        where F: for<'a> FnOnce(TyCtxt<'a, 'gcx, 'tcx>) -> R
+    pub fn enter_local<F, R>(
+        &self,
+        arena: &'tcx DroplessArena,
+        f: F
+    ) -> R
+    where
+        F: for<'a> FnOnce(TyCtxt<'a, 'gcx, 'tcx>) -> R
     {
         let interners = CtxtInterners::new(arena);
-        tls::enter(self, &interners, f)
+        let tcx = TyCtxt {
+            gcx: self,
+            interners: &interners,
+        };
+        ty::tls::with_related_context(tcx.global_tcx(), |icx| {
+            let new_icx = ty::tls::ImplicitCtxt {
+                tcx,
+                query: icx.query.clone(),
+            };
+            ty::tls::enter_context(&new_icx, |new_icx| {
+                f(new_icx.tcx)
+            })
+        })
     }
 }
 
@@ -1678,83 +1697,196 @@ impl<'a, 'tcx> Lift<'tcx> for &'a Slice<CanonicalVarInfo> {
 }
 
 pub mod tls {
-    use super::{CtxtInterners, GlobalCtxt, TyCtxt};
+    use super::{GlobalCtxt, TyCtxt};
 
     use std::cell::Cell;
     use std::fmt;
+    use std::mem;
     use syntax_pos;
+    use ty::maps;
+    use errors::{Diagnostic, TRACK_DIAGNOSTICS};
+    use rustc_data_structures::OnDrop;
+    use rustc_data_structures::sync::Lrc;
 
-    /// Marker types used for the scoped TLS slot.
-    /// The type context cannot be used directly because the scoped TLS
-    /// in libstd doesn't allow types generic over lifetimes.
-    enum ThreadLocalGlobalCtxt {}
-    enum ThreadLocalInterners {}
+    /// This is the implicit state of rustc. It contains the current
+    /// TyCtxt and query. It is updated when creating a local interner or
+    /// executing a new query. Whenever there's a TyCtxt value available
+    /// you should also have access to an ImplicitCtxt through the functions
+    /// in this module.
+    #[derive(Clone)]
+    pub struct ImplicitCtxt<'a, 'gcx: 'a+'tcx, 'tcx: 'a> {
+        /// The current TyCtxt. Initially created by `enter_global` and updated
+        /// by `enter_local` with a new local interner
+        pub tcx: TyCtxt<'a, 'gcx, 'tcx>,
 
-    thread_local! {
-        static TLS_TCX: Cell<Option<(*const ThreadLocalGlobalCtxt,
-                                     *const ThreadLocalInterners)>> = Cell::new(None)
+        /// The current query job, if any. This is updated by start_job in
+        /// ty::maps::plumbing when executing a query
+        pub query: Option<Lrc<maps::QueryJob<'gcx>>>,
     }
 
+    // A thread local value which stores a pointer to the current ImplicitCtxt
+    thread_local!(static TLV: Cell<usize> = Cell::new(0));
+
+    fn set_tlv<F: FnOnce() -> R, R>(value: usize, f: F) -> R {
+        let old = get_tlv();
+        let _reset = OnDrop(move || TLV.with(|tlv| tlv.set(old)));
+        TLV.with(|tlv| tlv.set(value));
+        f()
+    }
+
+    fn get_tlv() -> usize {
+        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(|tcx| {
             write!(f, "{}", tcx.sess.codemap().span_to_string(span))
         })
     }
 
-    pub fn enter_global<'gcx, F, R>(gcx: GlobalCtxt<'gcx>, f: F) -> R
-        where F: for<'a> FnOnce(TyCtxt<'a, 'gcx, 'gcx>) -> R
+    /// 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(|context| {
+            if let Some(ref query) = context.query {
+                query.diagnostics.lock().push(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 result = enter(&gcx, &gcx.global_interners, f);
-            span_dbg.set(original_span_debug);
-            result
+
+            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()
+            })
         })
     }
 
-    pub fn enter<'a, 'gcx: 'tcx, 'tcx, F, R>(gcx: &'a GlobalCtxt<'gcx>,
-                                             interners: &'a CtxtInterners<'tcx>,
-                                             f: F) -> R
-        where F: FnOnce(TyCtxt<'a, 'gcx, 'tcx>) -> R
+    /// Sets `context` as the new current ImplicitCtxt for the duration of the function `f`
+    pub fn enter_context<'a, 'gcx: 'tcx, 'tcx, F, R>(context: &ImplicitCtxt<'a, 'gcx, 'tcx>,
+                                                     f: F) -> R
+        where F: FnOnce(&ImplicitCtxt<'a, 'gcx, 'tcx>) -> R
     {
-        let gcx_ptr = gcx as *const _ as *const ThreadLocalGlobalCtxt;
-        let interners_ptr = interners as *const _ as *const ThreadLocalInterners;
-        TLS_TCX.with(|tls| {
-            let prev = tls.get();
-            tls.set(Some((gcx_ptr, interners_ptr)));
-            let ret = f(TyCtxt {
-                gcx,
-                interners,
-            });
-            tls.set(prev);
-            ret
+        set_tlv(context as *const _ as usize, || {
+            f(&context)
         })
     }
 
-    pub fn with<F, R>(f: F) -> R
-        where F: for<'a, 'gcx, 'tcx> FnOnce(TyCtxt<'a, 'gcx, 'tcx>) -> R
+    /// Enters GlobalCtxt by setting up libsyntax callbacks and
+    /// creating a initial TyCtxt and ImplicitCtxt.
+    /// This happens once per rustc session and TyCtxts only exists
+    /// inside the `f` function.
+    pub fn enter_global<'gcx, F, R>(gcx: &GlobalCtxt<'gcx>, f: F) -> R
+        where F: for<'a> FnOnce(TyCtxt<'a, 'gcx, 'gcx>) -> R
     {
-        TLS_TCX.with(|tcx| {
-            let (gcx, interners) = tcx.get().unwrap();
-            let gcx = unsafe { &*(gcx as *const GlobalCtxt) };
-            let interners = unsafe { &*(interners as *const CtxtInterners) };
-            f(TyCtxt {
+        with_thread_locals(|| {
+            let tcx = TyCtxt {
                 gcx,
-                interners,
+                interners: &gcx.global_interners,
+            };
+            let icx = ImplicitCtxt {
+                tcx,
+                query: None,
+            };
+            enter_context(&icx, |_| {
+                f(tcx)
             })
         })
     }
 
-    pub fn with_opt<F, R>(f: F) -> R
-        where F: for<'a, 'gcx, 'tcx> FnOnce(Option<TyCtxt<'a, 'gcx, 'tcx>>) -> R
+    /// Allows access to the current ImplicitCtxt in a closure if one is available
+    pub fn with_context_opt<F, R>(f: F) -> R
+        where F: for<'a, 'gcx, 'tcx> FnOnce(Option<&ImplicitCtxt<'a, 'gcx, 'tcx>>) -> R
     {
-        if TLS_TCX.with(|tcx| tcx.get().is_some()) {
-            with(|v| f(Some(v)))
-        } else {
+        let context = get_tlv();
+        if context == 0 {
             f(None)
+        } else {
+            unsafe { f(Some(&*(context as *const ImplicitCtxt))) }
         }
     }
+
+    /// Allows access to the current ImplicitCtxt.
+    /// Panics if there is no ImplicitCtxt available
+    pub fn with_context<F, R>(f: F) -> R
+        where F: for<'a, 'gcx, 'tcx> FnOnce(&ImplicitCtxt<'a, 'gcx, 'tcx>) -> R
+    {
+        with_context_opt(|opt_context| f(opt_context.expect("no ImplicitCtxt stored in tls")))
+    }
+
+    /// Allows access to the current ImplicitCtxt whose tcx field has the same global
+    /// interner as the tcx argument passed in. This means the closure is given an ImplicitCtxt
+    /// with the same 'gcx lifetime as the TyCtxt passed in.
+    /// This will panic if you pass it a TyCtxt which has a different global interner from
+    /// the current ImplicitCtxt's tcx field.
+    pub fn with_related_context<'a, 'gcx, 'tcx1, F, R>(tcx: TyCtxt<'a, 'gcx, 'tcx1>, f: F) -> R
+        where F: for<'b, 'tcx2> FnOnce(&ImplicitCtxt<'b, 'gcx, 'tcx2>) -> R
+    {
+        with_context(|context| {
+            unsafe {
+                let gcx = tcx.gcx as *const _ as usize;
+                assert!(context.tcx.gcx as *const _ as usize == gcx);
+                let context: &ImplicitCtxt = mem::transmute(context);
+                f(context)
+            }
+        })
+    }
+
+    /// Allows access to the current ImplicitCtxt whose tcx field has the same global
+    /// interner and local interner as the tcx argument passed in. This means the closure
+    /// is given an ImplicitCtxt with the same 'tcx and 'gcx lifetimes as the TyCtxt passed in.
+    /// This will panic if you pass it a TyCtxt which has a different global interner or
+    /// a different local interner from the current ImplicitCtxt's tcx field.
+    pub fn with_fully_related_context<'a, 'gcx, 'tcx, F, R>(tcx: TyCtxt<'a, 'gcx, 'tcx>, f: F) -> R
+        where F: for<'b> FnOnce(&ImplicitCtxt<'b, 'gcx, 'tcx>) -> R
+    {
+        with_context(|context| {
+            unsafe {
+                let gcx = tcx.gcx as *const _ as usize;
+                let interners = tcx.interners as *const _ as usize;
+                assert!(context.tcx.gcx as *const _ as usize == gcx);
+                assert!(context.tcx.interners as *const _ as usize == interners);
+                let context: &ImplicitCtxt = mem::transmute(context);
+                f(context)
+            }
+        })
+    }
+
+    /// Allows access to the TyCtxt in the current ImplicitCtxt.
+    /// Panics if there is no ImplicitCtxt available
+    pub fn with<F, R>(f: F) -> R
+        where F: for<'a, 'gcx, 'tcx> FnOnce(TyCtxt<'a, 'gcx, 'tcx>) -> R
+    {
+        with_context(|context| f(context.tcx))
+    }
+
+    /// Allows access to the TyCtxt in the current ImplicitCtxt.
+    /// The closure is passed None if there is no ImplicitCtxt available
+    pub fn with_opt<F, R>(f: F) -> R
+        where F: for<'a, 'gcx, 'tcx> FnOnce(Option<TyCtxt<'a, 'gcx, 'tcx>>) -> R
+    {
+        with_context_opt(|opt_context| f(opt_context.map(|context| context.tcx)))
+    }
 }
 
 macro_rules! sty_debug_print {

src/librustc/ty/maps/job.rs

@@ -0,0 +1,90 @@
+// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use rustc_data_structures::sync::{Lock, Lrc};
+use syntax_pos::Span;
+use ty::tls;
+use ty::maps::Query;
+use ty::maps::plumbing::CycleError;
+use ty::context::TyCtxt;
+use errors::Diagnostic;
+
+/// Indicates the state of a query for a given key in a query map
+pub(super) enum QueryResult<'tcx, T> {
+    /// An already executing query. The query job can be used to await for its completion
+    Started(Lrc<QueryJob<'tcx>>),
+
+    /// The query is complete and produced `T`
+    Complete(T),
+
+    /// The query panicked. Queries trying to wait on this will raise a fatal error / silently panic
+    Poisoned,
+}
+
+/// A span and a query key
+#[derive(Clone, Debug)]
+pub struct QueryInfo<'tcx> {
+    pub span: Span,
+    pub query: Query<'tcx>,
+}
+
+/// A object representing an active query job.
+pub struct QueryJob<'tcx> {
+    pub info: QueryInfo<'tcx>,
+
+    /// The parent query job which created this job and is implicitly waiting on it.
+    pub parent: Option<Lrc<QueryJob<'tcx>>>,
+
+    /// Diagnostic messages which are emitted while the query executes
+    pub diagnostics: Lock<Vec<Diagnostic>>,
+}
+
+impl<'tcx> QueryJob<'tcx> {
+    /// Creates a new query job
+    pub fn new(info: QueryInfo<'tcx>, parent: Option<Lrc<QueryJob<'tcx>>>) -> Self {
+        QueryJob {
+            diagnostics: Lock::new(Vec::new()),
+            info,
+            parent,
+        }
+    }
+
+    /// Awaits for the query job to complete.
+    ///
+    /// For single threaded rustc there's no concurrent jobs running, so if we are waiting for any
+    /// query that means that there is a query cycle, thus this always running a cycle error.
+    pub(super) fn await<'lcx>(
+        &self,
+        tcx: TyCtxt<'_, 'tcx, 'lcx>,
+        span: Span,
+    ) -> Result<(), CycleError<'tcx>> {
+        // Get the current executing query (waiter) and find the waitee amongst its parents
+        let mut current_job = tls::with_related_context(tcx, |icx| icx.query.clone());
+        let mut cycle = Vec::new();
+
+        while let Some(job) = current_job {
+            cycle.insert(0, job.info.clone());
+
+            if &*job as *const _ == self as *const _ {
+                break;
+            }
+
+            current_job = job.parent.clone();
+        }
+
+        Err(CycleError { span, cycle })
+    }
+
+    /// Signals to waiters that the query is complete.
+    ///
+    /// This does nothing for single threaded rustc,
+    /// as there are no concurrent jobs which could be waiting on us
+    pub fn signal_complete(&self) {}
+}