Rework SESSION_GLOBALS API to prevent overwriting it

Author: GuillaumeGomez

compiler/rustc_ast/src/util/comments/tests.rs

@@ -1,9 +1,9 @@
 use super::*;
-use rustc_span::with_default_session_globals;
+use rustc_span::create_default_session_globals_then;
 
 #[test]
 fn test_block_doc_comment_1() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let comment = "\n * Test \n **  Test\n *   Test\n";
         let stripped = beautify_doc_string(Symbol::intern(comment));
         assert_eq!(stripped.as_str(), " Test \n*  Test\n   Test");
@@ -12,7 +12,7 @@ fn test_block_doc_comment_1() {
 
 #[test]
 fn test_block_doc_comment_2() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let comment = "\n * Test\n *  Test\n";
         let stripped = beautify_doc_string(Symbol::intern(comment));
         assert_eq!(stripped.as_str(), " Test\n  Test");
@@ -21,7 +21,7 @@ fn test_block_doc_comment_2() {
 
 #[test]
 fn test_block_doc_comment_3() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let comment = "\n let a: *i32;\n *a = 5;\n";
         let stripped = beautify_doc_string(Symbol::intern(comment));
         assert_eq!(stripped.as_str(), " let a: *i32;\n *a = 5;");
@@ -30,7 +30,7 @@ fn test_block_doc_comment_3() {
 
 #[test]
 fn test_line_doc_comment() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let stripped = beautify_doc_string(Symbol::intern(" test"));
         assert_eq!(stripped.as_str(), " test");
         let stripped = beautify_doc_string(Symbol::intern("! test"));

compiler/rustc_ast_pretty/src/pprust/tests.rs

@@ -1,8 +1,8 @@
 use super::*;
 
 use rustc_ast as ast;
+use rustc_span::create_default_session_globals_then;
 use rustc_span::symbol::Ident;
-use rustc_span::with_default_session_globals;
 
 fn fun_to_string(
     decl: &ast::FnDecl,
@@ -24,7 +24,7 @@ fn variant_to_string(var: &ast::Variant) -> String {
 
 #[test]
 fn test_fun_to_string() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let abba_ident = Ident::from_str("abba");
 
         let decl =
@@ -39,7 +39,7 @@ fn test_fun_to_string() {
 
 #[test]
 fn test_variant_to_string() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let ident = Ident::from_str("principal_skinner");
 
         let var = ast::Variant {

compiler/rustc_errors/src/json/tests.rs

@@ -39,16 +39,11 @@ impl<T: Write> Write for Shared<T> {
     }
 }
 
-fn with_default_session_globals(f: impl FnOnce()) {
-    let session_globals = rustc_span::SessionGlobals::new(rustc_span::edition::DEFAULT_EDITION);
-    rustc_span::SESSION_GLOBALS.set(&session_globals, f);
-}
-
 /// Test the span yields correct positions in JSON.
 fn test_positions(code: &str, span: (u32, u32), expected_output: SpanTestData) {
     let expected_output = TestData { spans: vec![expected_output] };
 
-    with_default_session_globals(|| {
+    rustc_span::create_default_session_globals_then(|| {
         let sm = Lrc::new(SourceMap::new(FilePathMapping::empty()));
         sm.new_source_file(Path::new("test.rs").to_owned().into(), code.to_owned());
 

compiler/rustc_expand/src/mut_visit/tests.rs

@@ -3,8 +3,8 @@ use crate::tests::{matches_codepattern, string_to_crate};
 use rustc_ast as ast;
 use rustc_ast::mut_visit::MutVisitor;
 use rustc_ast_pretty::pprust;
+use rustc_span::create_default_session_globals_then;
 use rustc_span::symbol::Ident;
-use rustc_span::with_default_session_globals;
 
 // This version doesn't care about getting comments or doc-strings in.
 fn print_crate_items(krate: &ast::Crate) -> String {
@@ -38,7 +38,7 @@ macro_rules! assert_pred {
 // Make sure idents get transformed everywhere.
 #[test]
 fn ident_transformation() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let mut zz_visitor = ToZzIdentMutVisitor;
         let mut krate =
             string_to_crate("#[a] mod b {fn c (d : e, f : g) {h!(i,j,k);l;m}}".to_string());
@@ -55,7 +55,7 @@ fn ident_transformation() {
 // Make sure idents get transformed even inside macro defs.
 #[test]
 fn ident_transformation_in_defs() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let mut zz_visitor = ToZzIdentMutVisitor;
         let mut krate = string_to_crate(
             "macro_rules! a {(b $c:expr $(d $e:token)f+ => \

compiler/rustc_expand/src/parse/tests.rs

@@ -10,9 +10,9 @@ use rustc_errors::PResult;
 use rustc_parse::new_parser_from_source_str;
 use rustc_parse::parser::ForceCollect;
 use rustc_session::parse::ParseSess;
+use rustc_span::create_default_session_globals_then;
 use rustc_span::source_map::FilePathMapping;
 use rustc_span::symbol::{kw, sym, Symbol};
-use rustc_span::with_default_session_globals;
 use rustc_span::{BytePos, FileName, Pos, Span};
 
 use std::path::PathBuf;
@@ -51,15 +51,15 @@ fn string_to_item(source_str: String) -> Option<P<ast::Item>> {
 #[should_panic]
 #[test]
 fn bad_path_expr_1() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         string_to_expr("::abc::def::return".to_string());
     })
 }
 
 // Checks the token-tree-ization of macros.
 #[test]
 fn string_to_tts_macro() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let tts: Vec<_> =
             string_to_stream("macro_rules! zip (($a)=>($a))".to_string()).trees().collect();
         let tts: &[TokenTree] = &tts[..];
@@ -96,7 +96,7 @@ fn string_to_tts_macro() {
 
 #[test]
 fn string_to_tts_1() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let tts = string_to_stream("fn a (b : i32) { b; }".to_string());
 
         let expected = TokenStream::new(vec![
@@ -131,7 +131,7 @@ fn string_to_tts_1() {
 
 #[test]
 fn parse_use() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let use_s = "use foo::bar::baz;";
         let vitem = string_to_item(use_s.to_string()).unwrap();
         let vitem_s = item_to_string(&vitem);
@@ -146,7 +146,7 @@ fn parse_use() {
 
 #[test]
 fn parse_extern_crate() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let ex_s = "extern crate foo;";
         let vitem = string_to_item(ex_s.to_string()).unwrap();
         let vitem_s = item_to_string(&vitem);
@@ -184,7 +184,7 @@ fn get_spans_of_pat_idents(src: &str) -> Vec<Span> {
 
 #[test]
 fn span_of_self_arg_pat_idents_are_correct() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let srcs = [
             "impl z { fn a (&self, &myarg: i32) {} }",
             "impl z { fn a (&mut self, &myarg: i32) {} }",
@@ -208,7 +208,7 @@ fn span_of_self_arg_pat_idents_are_correct() {
 
 #[test]
 fn parse_exprs() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         // just make sure that they parse....
         string_to_expr("3 + 4".to_string());
         string_to_expr("a::z.froob(b,&(987+3))".to_string());
@@ -217,7 +217,7 @@ fn parse_exprs() {
 
 #[test]
 fn attrs_fix_bug() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         string_to_item(
             "pub fn mk_file_writer(path: &Path, flags: &[FileFlag])
                 -> Result<Box<Writer>, String> {
@@ -238,7 +238,7 @@ let mut fflags: c_int = wb();
 
 #[test]
 fn crlf_doc_comments() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let sess = sess();
 
         let name_1 = FileName::Custom("crlf_source_1".to_string());
@@ -272,7 +272,7 @@ fn ttdelim_span() {
         new_parser_from_source_str(sess, name, source).parse_expr()
     }
 
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let sess = sess();
         let expr = parse_expr_from_source_str(
             PathBuf::from("foo").into(),
@@ -300,7 +300,7 @@ fn ttdelim_span() {
 // See `recurse_into_file_modules` in the parser.
 #[test]
 fn out_of_line_mod() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let item = parse_item_from_source_str(
             PathBuf::from("foo").into(),
             "mod foo { struct S; mod this_does_not_exist; }".to_owned(),

compiler/rustc_expand/src/tests.rs

@@ -2,8 +2,8 @@ use rustc_ast as ast;
 use rustc_ast::tokenstream::TokenStream;
 use rustc_parse::{new_parser_from_source_str, parser::Parser, source_file_to_stream};
 use rustc_session::parse::ParseSess;
+use rustc_span::create_default_session_if_not_set_then;
 use rustc_span::source_map::{FilePathMapping, SourceMap};
-use rustc_span::with_default_session_globals;
 use rustc_span::{BytePos, MultiSpan, Span};
 
 use rustc_data_structures::sync::Lrc;
@@ -124,7 +124,7 @@ impl<T: Write> Write for Shared<T> {
 }
 
 fn test_harness(file_text: &str, span_labels: Vec<SpanLabel>, expected_output: &str) {
-    with_default_session_globals(|| {
+    create_default_session_if_not_set_then(|_| {
         let output = Arc::new(Mutex::new(Vec::new()));
 
         let source_map = Lrc::new(SourceMap::new(FilePathMapping::empty()));

compiler/rustc_expand/src/tokenstream/tests.rs

@@ -2,7 +2,7 @@ use crate::tests::string_to_stream;
 
 use rustc_ast::token;
 use rustc_ast::tokenstream::{Spacing, TokenStream, TokenStreamBuilder, TokenTree};
-use rustc_span::with_default_session_globals;
+use rustc_span::create_default_session_globals_then;
 use rustc_span::{BytePos, Span, Symbol};
 use smallvec::smallvec;
 
@@ -20,7 +20,7 @@ fn joint(tree: TokenTree) -> TokenStream {
 
 #[test]
 fn test_concat() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let test_res = string_to_ts("foo::bar::baz");
         let test_fst = string_to_ts("foo::bar");
         let test_snd = string_to_ts("::baz");
@@ -33,7 +33,7 @@ fn test_concat() {
 
 #[test]
 fn test_to_from_bijection() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let test_start = string_to_ts("foo::bar(baz)");
         let test_end = test_start.trees().collect();
         assert_eq!(test_start, test_end)
@@ -42,7 +42,7 @@ fn test_to_from_bijection() {
 
 #[test]
 fn test_eq_0() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let test_res = string_to_ts("foo");
         let test_eqs = string_to_ts("foo");
         assert_eq!(test_res, test_eqs)
@@ -51,7 +51,7 @@ fn test_eq_0() {
 
 #[test]
 fn test_eq_1() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let test_res = string_to_ts("::bar::baz");
         let test_eqs = string_to_ts("::bar::baz");
         assert_eq!(test_res, test_eqs)
@@ -60,7 +60,7 @@ fn test_eq_1() {
 
 #[test]
 fn test_eq_3() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let test_res = string_to_ts("");
         let test_eqs = string_to_ts("");
         assert_eq!(test_res, test_eqs)
@@ -69,7 +69,7 @@ fn test_eq_3() {
 
 #[test]
 fn test_diseq_0() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let test_res = string_to_ts("::bar::baz");
         let test_eqs = string_to_ts("bar::baz");
         assert_eq!(test_res == test_eqs, false)
@@ -78,7 +78,7 @@ fn test_diseq_0() {
 
 #[test]
 fn test_diseq_1() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let test_res = string_to_ts("(bar,baz)");
         let test_eqs = string_to_ts("bar,baz");
         assert_eq!(test_res == test_eqs, false)
@@ -87,7 +87,7 @@ fn test_diseq_1() {
 
 #[test]
 fn test_is_empty() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let test0: TokenStream = Vec::<TokenTree>::new().into_iter().collect();
         let test1: TokenStream =
             TokenTree::token(token::Ident(Symbol::intern("a"), false), sp(0, 1)).into();
@@ -101,7 +101,7 @@ fn test_is_empty() {
 
 #[test]
 fn test_dotdotdot() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let mut builder = TokenStreamBuilder::new();
         builder.push(joint(TokenTree::token(token::Dot, sp(0, 1))));
         builder.push(joint(TokenTree::token(token::Dot, sp(1, 2))));

compiler/rustc_interface/src/interface.rs

@@ -76,7 +76,7 @@ impl Compiler {
 
 /// Converts strings provided as `--cfg [cfgspec]` into a `crate_cfg`.
 pub fn parse_cfgspecs(cfgspecs: Vec<String>) -> FxHashSet<(String, Option<String>)> {
-    rustc_span::with_default_session_globals(move || {
+    rustc_span::create_default_session_if_not_set_then(move |_| {
         let cfg = cfgspecs
             .into_iter()
             .map(|s| {

compiler/rustc_interface/src/tests.rs

@@ -99,7 +99,7 @@ fn assert_different_hash(x: &Options, y: &Options) {
 // When the user supplies --test we should implicitly supply --cfg test
 #[test]
 fn test_switch_implies_cfg_test() {
-    rustc_span::with_default_session_globals(|| {
+    rustc_span::create_default_session_globals_then(|| {
         let matches = optgroups().parse(&["--test".to_string()]).unwrap();
         let (sess, cfg) = mk_session(matches);
         let cfg = build_configuration(&sess, to_crate_config(cfg));
@@ -110,7 +110,7 @@ fn test_switch_implies_cfg_test() {
 // When the user supplies --test and --cfg test, don't implicitly add another --cfg test
 #[test]
 fn test_switch_implies_cfg_test_unless_cfg_test() {
-    rustc_span::with_default_session_globals(|| {
+    rustc_span::create_default_session_globals_then(|| {
         let matches = optgroups().parse(&["--test".to_string(), "--cfg=test".to_string()]).unwrap();
         let (sess, cfg) = mk_session(matches);
         let cfg = build_configuration(&sess, to_crate_config(cfg));
@@ -122,20 +122,20 @@ fn test_switch_implies_cfg_test_unless_cfg_test() {
 
 #[test]
 fn test_can_print_warnings() {
-    rustc_span::with_default_session_globals(|| {
+    rustc_span::create_default_session_globals_then(|| {
         let matches = optgroups().parse(&["-Awarnings".to_string()]).unwrap();
         let (sess, _) = mk_session(matches);
         assert!(!sess.diagnostic().can_emit_warnings());
     });
 
-    rustc_span::with_default_session_globals(|| {
+    rustc_span::create_default_session_globals_then(|| {
         let matches =
             optgroups().parse(&["-Awarnings".to_string(), "-Dwarnings".to_string()]).unwrap();
         let (sess, _) = mk_session(matches);
         assert!(sess.diagnostic().can_emit_warnings());
     });
 
-    rustc_span::with_default_session_globals(|| {
+    rustc_span::create_default_session_globals_then(|| {
         let matches = optgroups().parse(&["-Adead_code".to_string()]).unwrap();
         let (sess, _) = mk_session(matches);
         assert!(sess.diagnostic().can_emit_warnings());

compiler/rustc_interface/src/util.rs

@@ -153,7 +153,7 @@ pub fn setup_callbacks_and_run_in_thread_pool_with_globals<F: FnOnce() -> R + Se
     crate::callbacks::setup_callbacks();
 
     let main_handler = move || {
-        rustc_span::with_session_globals(edition, || {
+        rustc_span::create_session_globals_then(edition, || {
             io::set_output_capture(stderr.clone());
             f()
         })
@@ -174,12 +174,13 @@ unsafe fn handle_deadlock() {
     rustc_data_structures::sync::assert_sync::<tls::ImplicitCtxt<'_, '_>>();
     let icx: &tls::ImplicitCtxt<'_, '_> = &*(context as *const tls::ImplicitCtxt<'_, '_>);
 
-    let session_globals = rustc_span::SESSION_GLOBALS.with(|sg| sg as *const _);
+    let session_globals = rustc_span::with_session_globals(|sg| sg as *const _);
     let session_globals = &*session_globals;
     thread::spawn(move || {
         tls::enter_context(icx, |_| {
-            rustc_span::SESSION_GLOBALS
-                .set(session_globals, || tls::with(|tcx| tcx.queries.deadlock(tcx, &registry)))
+            rustc_span::set_session_globals_then(session_globals, || {
+                tls::with(|tcx| tcx.queries.deadlock(tcx, &registry))
+            })
         });
     });
 }
@@ -206,13 +207,13 @@ pub fn setup_callbacks_and_run_in_thread_pool_with_globals<F: FnOnce() -> R + Se
 
     let with_pool = move |pool: &rayon::ThreadPool| pool.install(f);
 
-    rustc_span::with_session_globals(edition, || {
-        rustc_span::SESSION_GLOBALS.with(|session_globals| {
+    rustc_span::create_session_globals_then(edition, || {
+        rustc_span::with_session_globals(|session_globals| {
             // The main handler runs for each Rayon worker thread and sets up
             // the thread local rustc uses. `session_globals` is captured and set
             // on the new threads.
             let main_handler = move |thread: rayon::ThreadBuilder| {
-                rustc_span::SESSION_GLOBALS.set(session_globals, || {
+                rustc_span::set_session_globals_then(session_globals, || {
                     io::set_output_capture(stderr.clone());
                     thread.run()
                 })

compiler/rustc_mir/src/transform/coverage/tests.rs

@@ -22,7 +22,7 @@
 //! are unrelated to the `TyCtxt` global. Without initializing the `Span` session globals, some
 //! basic, coverage-specific features would be impossible to test, but thankfully initializing these
 //! globals is comparatively simpler. The easiest way is to wrap the test in a closure argument
-//! to: `rustc_span::with_default_session_globals(|| { test_here(); })`.
+//! to: `rustc_span::create_default_session_globals_then(|| { test_here(); })`.
 
 use super::counters;
 use super::debug;
@@ -677,7 +677,7 @@ fn synthesize_body_span_from_terminators(mir_body: &Body<'_>) -> Span {
 
 #[test]
 fn test_make_bcb_counters() {
-    rustc_span::with_default_session_globals(|| {
+    rustc_span::create_default_session_globals_then(|| {
         let mir_body = goto_switchint();
         let body_span = synthesize_body_span_from_terminators(&mir_body);
         let mut basic_coverage_blocks = graph::CoverageGraph::from_mir(&mir_body);

compiler/rustc_parse_format/src/tests.rs

@@ -144,8 +144,7 @@ fn format_align_fill() {
 }
 #[test]
 fn format_counts() {
-    use rustc_span::{edition, SessionGlobals, SESSION_GLOBALS};
-    SESSION_GLOBALS.set(&SessionGlobals::new(edition::DEFAULT_EDITION), || {
+    rustc_span::create_default_session_globals_then(|| {
         same(
             "{:10x}",
             &[NextArgument(Argument {

compiler/rustc_span/src/hygiene.rs

@@ -26,7 +26,7 @@
 
 use crate::edition::Edition;
 use crate::symbol::{kw, sym, Symbol};
-use crate::SESSION_GLOBALS;
+use crate::with_session_globals;
 use crate::{BytePos, CachingSourceMapView, ExpnIdCache, SourceFile, Span, DUMMY_SP};
 
 use crate::def_id::{CrateNum, DefId, CRATE_DEF_INDEX, LOCAL_CRATE};
@@ -200,7 +200,7 @@ impl HygieneData {
     }
 
     pub fn with<T, F: FnOnce(&mut HygieneData) -> T>(f: F) -> T {
-        SESSION_GLOBALS.with(|session_globals| f(&mut *session_globals.hygiene_data.borrow_mut()))
+        with_session_globals(|session_globals| f(&mut *session_globals.hygiene_data.borrow_mut()))
     }
 
     fn fresh_expn(&mut self, mut expn_data: Option<ExpnData>) -> ExpnId {
@@ -1359,8 +1359,9 @@ fn update_disambiguator(expn_id: ExpnId) {
         }
     }
 
-    let source_map = SESSION_GLOBALS
-        .with(|session_globals| session_globals.source_map.borrow().as_ref().unwrap().clone());
+    let source_map = with_session_globals(|session_globals| {
+        session_globals.source_map.borrow().as_ref().unwrap().clone()
+    });
 
     let mut ctx =
         DummyHashStableContext { caching_source_map: CachingSourceMapView::new(&source_map) };

compiler/rustc_span/src/lev_distance/tests.rs

@@ -21,8 +21,8 @@ fn test_lev_distance() {
 
 #[test]
 fn test_find_best_match_for_name() {
-    use crate::with_default_session_globals;
-    with_default_session_globals(|| {
+    use crate::create_default_session_globals_then;
+    create_default_session_globals_then(|| {
         let input = vec![Symbol::intern("aaab"), Symbol::intern("aaabc")];
         assert_eq!(
             find_best_match_for_name(&input, Symbol::intern("aaaa"), None),

compiler/rustc_span/src/lib.rs

@@ -99,19 +99,65 @@ impl SessionGlobals {
     }
 }
 
-pub fn with_session_globals<R>(edition: Edition, f: impl FnOnce() -> R) -> R {
+#[inline]
+pub fn create_session_globals_then<R>(edition: Edition, f: impl FnOnce() -> R) -> R {
+    assert!(
+        !SESSION_GLOBALS.is_set(),
+        "SESSION_GLOBALS should never be overwritten! \
+         Use another thread if you need another SessionGlobals"
+    );
     let session_globals = SessionGlobals::new(edition);
     SESSION_GLOBALS.set(&session_globals, f)
 }
 
-pub fn with_default_session_globals<R>(f: impl FnOnce() -> R) -> R {
-    with_session_globals(edition::DEFAULT_EDITION, f)
+#[inline]
+pub fn set_session_globals_then<R>(session_globals: &SessionGlobals, f: impl FnOnce() -> R) -> R {
+    assert!(
+        !SESSION_GLOBALS.is_set(),
+        "SESSION_GLOBALS should never be overwritten! \
+         Use another thread if you need another SessionGlobals"
+    );
+    SESSION_GLOBALS.set(session_globals, f)
+}
+
+#[inline]
+pub fn create_default_session_if_not_set_then<R, F>(f: F) -> R
+where
+    F: FnOnce(&SessionGlobals) -> R,
+{
+    create_session_if_not_set_then(edition::DEFAULT_EDITION, f)
+}
+
+#[inline]
+pub fn create_session_if_not_set_then<R, F>(edition: Edition, f: F) -> R
+where
+    F: FnOnce(&SessionGlobals) -> R,
+{
+    if !SESSION_GLOBALS.is_set() {
+        let session_globals = SessionGlobals::new(edition);
+        SESSION_GLOBALS.set(&session_globals, || SESSION_GLOBALS.with(f))
+    } else {
+        SESSION_GLOBALS.with(f)
+    }
+}
+
+#[inline]
+pub fn with_session_globals<R, F>(f: F) -> R
+where
+    F: FnOnce(&SessionGlobals) -> R,
+{
+    SESSION_GLOBALS.with(f)
+}
+
+#[inline]
+pub fn create_default_session_globals_then<R>(f: impl FnOnce() -> R) -> R {
+    create_session_globals_then(edition::DEFAULT_EDITION, f)
 }
 
 // If this ever becomes non thread-local, `decode_syntax_context`
 // and `decode_expn_id` will need to be updated to handle concurrent
 // deserialization.
-scoped_tls::scoped_thread_local!(pub static SESSION_GLOBALS: SessionGlobals);
+scoped_tls::scoped_thread_local!(static SESSION_GLOBALS: SessionGlobals);
 
 // FIXME: We should use this enum or something like it to get rid of the
 // use of magic `/rust/1.x/...` paths across the board.
@@ -857,13 +903,13 @@ impl<D: Decoder> Decodable<D> for Span {
 /// the `SourceMap` provided to this function. If that is not available,
 /// we fall back to printing the raw `Span` field values.
 pub fn with_source_map<T, F: FnOnce() -> T>(source_map: Lrc<SourceMap>, f: F) -> T {
-    SESSION_GLOBALS.with(|session_globals| {
+    with_session_globals(|session_globals| {
         *session_globals.source_map.borrow_mut() = Some(source_map);
     });
     struct ClearSourceMap;
     impl Drop for ClearSourceMap {
         fn drop(&mut self) {
-            SESSION_GLOBALS.with(|session_globals| {
+            with_session_globals(|session_globals| {
                 session_globals.source_map.borrow_mut().take();
             });
         }
@@ -882,7 +928,7 @@ pub fn debug_with_source_map(
 }
 
 pub fn default_span_debug(span: Span, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-    SESSION_GLOBALS.with(|session_globals| {
+    with_session_globals(|session_globals| {
         if let Some(source_map) = &*session_globals.source_map.borrow() {
             debug_with_source_map(span, f, source_map)
         } else {

compiler/rustc_span/src/span_encoding.rs

@@ -5,7 +5,6 @@
 // See https://internals.rust-lang.org/t/rfc-compiler-refactoring-spans/1357/28
 
 use crate::hygiene::SyntaxContext;
-use crate::SESSION_GLOBALS;
 use crate::{BytePos, SpanData};
 
 use rustc_data_structures::fx::FxIndexSet;
@@ -122,5 +121,5 @@ impl SpanInterner {
 // If an interner exists, return it. Otherwise, prepare a fresh one.
 #[inline]
 fn with_span_interner<T, F: FnOnce(&mut SpanInterner) -> T>(f: F) -> T {
-    SESSION_GLOBALS.with(|session_globals| f(&mut *session_globals.span_interner.lock()))
+    crate::with_session_globals(|session_globals| f(&mut *session_globals.span_interner.lock()))
 }

compiler/rustc_span/src/symbol.rs

@@ -13,7 +13,7 @@ use std::fmt;
 use std::hash::{Hash, Hasher};
 use std::str;
 
-use crate::{Edition, Span, DUMMY_SP, SESSION_GLOBALS};
+use crate::{with_session_globals, Edition, Span, DUMMY_SP};
 
 #[cfg(test)]
 mod tests;
@@ -1783,7 +1783,7 @@ impl Ident {
 
 #[inline]
 fn with_interner<T, F: FnOnce(&mut Interner) -> T>(f: F) -> T {
-    SESSION_GLOBALS.with(|session_globals| f(&mut *session_globals.symbol_interner.lock()))
+    with_session_globals(|session_globals| f(&mut *session_globals.symbol_interner.lock()))
 }
 
 /// An alternative to [`Symbol`], useful when the chars within the symbol need to

compiler/rustc_span/src/symbol/tests.rs

@@ -1,6 +1,6 @@
 use super::*;
 
-use crate::{edition, SessionGlobals};
+use crate::create_default_session_globals_then;
 
 #[test]
 fn interner_tests() {
@@ -18,7 +18,7 @@ fn interner_tests() {
 
 #[test]
 fn without_first_quote_test() {
-    SESSION_GLOBALS.set(&SessionGlobals::new(edition::DEFAULT_EDITION), || {
+    create_default_session_globals_then(|| {
         let i = Ident::from_str("'break");
         assert_eq!(i.without_first_quote().name, kw::Break);
     });

src/librustdoc/clean/cfg/tests.rs

@@ -2,8 +2,8 @@ use super::*;
 
 use rustc_ast::attr;
 use rustc_ast::Path;
+use rustc_span::create_default_session_globals_then;
 use rustc_span::symbol::{Ident, Symbol};
-use rustc_span::with_default_session_globals;
 use rustc_span::DUMMY_SP;
 
 fn word_cfg(s: &str) -> Cfg {
@@ -52,7 +52,7 @@ macro_rules! dummy_meta_item_list {
 
 #[test]
 fn test_cfg_not() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         assert_eq!(!Cfg::False, Cfg::True);
         assert_eq!(!Cfg::True, Cfg::False);
         assert_eq!(!word_cfg("test"), Cfg::Not(Box::new(word_cfg("test"))));
@@ -70,7 +70,7 @@ fn test_cfg_not() {
 
 #[test]
 fn test_cfg_and() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let mut x = Cfg::False;
         x &= Cfg::True;
         assert_eq!(x, Cfg::False);
@@ -154,7 +154,7 @@ fn test_cfg_and() {
 
 #[test]
 fn test_cfg_or() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let mut x = Cfg::True;
         x |= Cfg::False;
         assert_eq!(x, Cfg::True);
@@ -238,7 +238,7 @@ fn test_cfg_or() {
 
 #[test]
 fn test_parse_ok() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let mi = dummy_meta_item_word("all");
         assert_eq!(Cfg::parse(&mi), Ok(word_cfg("all")));
 
@@ -271,7 +271,7 @@ fn test_parse_ok() {
 
 #[test]
 fn test_parse_err() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let mi = attr::mk_name_value_item(Ident::from_str("foo"), LitKind::Bool(false), DUMMY_SP);
         assert!(Cfg::parse(&mi).is_err());
 
@@ -303,7 +303,7 @@ fn test_parse_err() {
 
 #[test]
 fn test_render_short_html() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         assert_eq!(word_cfg("unix").render_short_html(), "Unix");
         assert_eq!(name_value_cfg("target_os", "macos").render_short_html(), "macOS");
         assert_eq!(name_value_cfg("target_pointer_width", "16").render_short_html(), "16-bit");
@@ -358,7 +358,7 @@ fn test_render_short_html() {
 
 #[test]
 fn test_render_long_html() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         assert_eq!(
             word_cfg("unix").render_long_html(),
             "This is supported on <strong>Unix</strong> only."
@@ -442,7 +442,7 @@ fn test_render_long_html() {
 fn test_simplify_with() {
     // This is a tiny subset of things that could be simplified, but it likely covers 90% of
     // real world usecases well.
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let foo = word_cfg("foo");
         let bar = word_cfg("bar");
         let baz = word_cfg("baz");

src/librustdoc/doctest.rs

@@ -514,7 +514,7 @@ crate fn make_test(
     // Uses librustc_ast to parse the doctest and find if there's a main fn and the extern
     // crate already is included.
     let result = rustc_driver::catch_fatal_errors(|| {
-        rustc_span::with_session_globals(edition, || {
+        rustc_span::create_session_if_not_set_then(edition, |_| {
             use rustc_errors::emitter::{Emitter, EmitterWriter};
             use rustc_errors::Handler;
             use rustc_parse::maybe_new_parser_from_source_str;

src/librustdoc/html/highlight/tests.rs

@@ -1,8 +1,8 @@
 use super::write_code;
 use crate::html::format::Buffer;
 use expect_test::expect_file;
+use rustc_span::create_default_session_globals_then;
 use rustc_span::edition::Edition;
-use rustc_span::with_default_session_globals;
 
 const STYLE: &str = r#"
 <style>
@@ -18,7 +18,7 @@ const STYLE: &str = r#"
 
 #[test]
 fn test_html_highlighting() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let src = include_str!("fixtures/sample.rs");
         let html = {
             let mut out = Buffer::new();
@@ -31,7 +31,7 @@ fn test_html_highlighting() {
 
 #[test]
 fn test_dos_backline() {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let src = "pub fn foo() {\r\n\
     println!(\"foo\");\r\n\
 }\r\n";

src/librustdoc/passes/unindent_comments/tests.rs

@@ -1,7 +1,7 @@
 use super::*;
+use rustc_span::create_default_session_globals_then;
 use rustc_span::source_map::DUMMY_SP;
 use rustc_span::symbol::Symbol;
-use rustc_span::with_default_session_globals;
 
 fn create_doc_fragment(s: &str) -> Vec<DocFragment> {
     vec![DocFragment {
@@ -17,7 +17,7 @@ fn create_doc_fragment(s: &str) -> Vec<DocFragment> {
 
 #[track_caller]
 fn run_test(input: &str, expected: &str) {
-    with_default_session_globals(|| {
+    create_default_session_globals_then(|| {
         let mut s = create_doc_fragment(input);
         unindent_fragments(&mut s);
         assert_eq!(&s.iter().collect::<String>(), expected);

src/test/ui-fulldeps/mod_dir_path_canonicalized.rs

@@ -19,7 +19,7 @@ use std::path::Path;
 mod gravy;
 
 pub fn main() {
-    rustc_span::with_default_session_globals(|| parse());
+    rustc_span::create_default_session_globals_then(|| parse());
 
     assert_eq!(gravy::foo(), 10);
 }

src/test/ui-fulldeps/pprust-expr-roundtrip.rs

@@ -212,7 +212,7 @@ impl MutVisitor for AddParens {
 }
 
 fn main() {
-    rustc_span::with_default_session_globals(|| run());
+    rustc_span::create_default_session_globals_then(|| run());
 }
 
 fn run() {

src/tools/clippy/clippy_lints/src/doc.rs

@@ -25,6 +25,7 @@ use rustc_span::source_map::{BytePos, FilePathMapping, MultiSpan, SourceMap, Spa
 use rustc_span::{sym, FileName, Pos};
 use std::io;
 use std::ops::Range;
+use std::thread;
 use url::Url;
 
 declare_clippy_lint! {
@@ -548,10 +549,10 @@ fn check_doc<'a, Events: Iterator<Item = (pulldown_cmark::Event<'a>, Range<usize
 }
 
 fn check_code(cx: &LateContext<'_>, text: &str, edition: Edition, span: Span) {
-    fn has_needless_main(code: &str, edition: Edition) -> bool {
+    fn has_needless_main(code: String, edition: Edition) -> bool {
         rustc_driver::catch_fatal_errors(|| {
-            rustc_span::with_session_globals(edition, || {
-                let filename = FileName::anon_source_code(code);
+            rustc_span::create_session_globals_then(edition, || {
+                let filename = FileName::anon_source_code(&code);
 
                 let sm = Lrc::new(SourceMap::new(FilePathMapping::empty()));
                 let emitter = EmitterWriter::new(box io::sink(), None, false, false, false, None, false);
@@ -613,7 +614,10 @@ fn check_code(cx: &LateContext<'_>, text: &str, edition: Edition, span: Span) {
         .unwrap_or_default()
     }
 
-    if has_needless_main(text, edition) {
+    // Because of the global session, we need to create a new session in a different thread with
+    // the edition we need.
+    let text = text.to_owned();
+    if thread::spawn(move || has_needless_main(text, edition)).join().expect("thread::spawn failed") {
         span_lint(cx, NEEDLESS_DOCTEST_MAIN, span, "needless `fn main` in doctest");
     }
 }

src/tools/error_index_generator/main.rs

@@ -283,7 +283,8 @@ fn parse_args() -> (OutputFormat, PathBuf) {
 fn main() {
     rustc_driver::init_env_logger("RUST_LOG");
     let (format, dst) = parse_args();
-    let result = rustc_span::with_default_session_globals(move || main_with_result(format, &dst));
+    let result =
+        rustc_span::create_default_session_globals_then(move || main_with_result(format, &dst));
     if let Err(e) = result {
         panic!("{}", e.to_string());
     }

src/tools/rustfmt/src/formatting.rs

@@ -34,7 +34,7 @@ impl<'b, T: Write + 'b> Session<'b, T> {
             return Err(ErrorKind::VersionMismatch);
         }
 
-        rustc_span::with_session_globals(self.config.edition().into(), || {
+        rustc_span::create_session_if_not_set_then(self.config.edition().into(), |_| {
             if self.config.disable_all_formatting() {
                 // When the input is from stdin, echo back the input.
                 if let Input::Text(ref buf) = input {