From 5357f83ee81bd1c11b4223f1531d9b2e2cb66ae7 Mon Sep 17 00:00:00 2001
From: Wesley Wiser <wwiser@gmail.com>
Date: Sun, 16 Feb 2020 20:32:25 -0500
Subject: [PATCH] [Miri] Use a session variable instead of checking for an env
 var always

In CTFE heavy code, checking the env var everytime is inefficient. We
can do a lot better by using a `Session` variable instead.
---
 src/librustc/mir/interpret/error.rs | 35 +++++++++++++++++------------
 src/librustc_session/session.rs     | 24 ++++++++++++++++++++
 2 files changed, 45 insertions(+), 14 deletions(-)

diff --git a/src/librustc/mir/interpret/error.rs b/src/librustc/mir/interpret/error.rs
index b46095927b7..0b33408edf0 100644
--- a/src/librustc/mir/interpret/error.rs
+++ b/src/librustc/mir/interpret/error.rs
@@ -5,15 +5,18 @@ use crate::mir;
 use crate::mir::interpret::ConstValue;
 use crate::ty::layout::{Align, LayoutError, Size};
 use crate::ty::query::TyCtxtAt;
+use crate::ty::tls;
 use crate::ty::{self, layout, Ty};
 
 use backtrace::Backtrace;
+use rustc_data_structures::sync::Lock;
 use rustc_errors::{struct_span_err, DiagnosticBuilder};
 use rustc_hir as hir;
 use rustc_macros::HashStable;
+use rustc_session::CtfeBacktrace;
 use rustc_span::{Pos, Span};
 use rustc_target::spec::abi::Abi;
-use std::{any::Any, env, fmt};
+use std::{any::Any, fmt};
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, HashStable, RustcEncodable, RustcDecodable)]
 pub enum ErrorHandled {
@@ -257,21 +260,25 @@ impl From<ErrorHandled> for InterpErrorInfo<'_> {
 
 impl<'tcx> From<InterpError<'tcx>> for InterpErrorInfo<'tcx> {
     fn from(kind: InterpError<'tcx>) -> Self {
-        let backtrace = match env::var("RUSTC_CTFE_BACKTRACE") {
-            // Matching `RUST_BACKTRACE` -- we treat "0" the same as "not present".
-            Ok(ref val) if val != "0" => {
-                let mut backtrace = Backtrace::new_unresolved();
-
-                if val == "immediate" {
-                    // Print it now.
-                    print_backtrace(&mut backtrace);
-                    None
-                } else {
-                    Some(Box::new(backtrace))
-                }
+        let capture_backtrace = tls::with_context_opt(|ctxt| {
+            if let Some(ctxt) = ctxt {
+                *Lock::borrow(&ctxt.tcx.sess.ctfe_backtrace)
+            } else {
+                CtfeBacktrace::Disabled
+            }
+        });
+
+        let backtrace = match capture_backtrace {
+            CtfeBacktrace::Disabled => None,
+            CtfeBacktrace::Capture => Some(Box::new(Backtrace::new_unresolved())),
+            CtfeBacktrace::Immediate => {
+                // Print it now.
+                let mut backtrace = Backtrace::new_unresolved();
+                print_backtrace(&mut backtrace);
+                None
             }
-            _ => None,
         };
+
         InterpErrorInfo { kind, backtrace }
     }
 }
diff --git a/src/librustc_session/session.rs b/src/librustc_session/session.rs
index 173b120e1f6..8cda95783a8 100644
--- a/src/librustc_session/session.rs
+++ b/src/librustc_session/session.rs
@@ -49,6 +49,18 @@ pub struct OptimizationFuel {
     out_of_fuel: bool,
 }
 
+/// The behavior of the CTFE engine when an error occurs with regards to backtraces.
+#[derive(Clone, Copy)]
+pub enum CtfeBacktrace {
+    /// Do nothing special, return the error as usual without a backtrace.
+    Disabled,
+    /// Capture a backtrace at the point the error is created and return it in the error
+    /// (to be printed later if/when the error ever actually gets shown to the user).
+    Capture,
+    /// Capture a backtrace at the point the error is created and immediately print it out.
+    Immediate,
+}
+
 /// Represents the data associated with a compilation
 /// session for a single crate.
 pub struct Session {
@@ -139,6 +151,11 @@ pub struct Session {
     /// Path for libraries that will take preference over libraries shipped by Rust.
     /// Used by windows-gnu targets to priortize system mingw-w64 libraries.
     pub system_library_path: OneThread<RefCell<Option<Option<PathBuf>>>>,
+
+    /// Tracks the current behavior of the CTFE engine when an error occurs.
+    /// Options range from returning the error without a backtrace to returning an error
+    /// and immediately printing the backtrace to stderr.
+    pub ctfe_backtrace: Lock<CtfeBacktrace>,
 }
 
 pub struct PerfStats {
@@ -1040,6 +1057,12 @@ fn build_session_(
         sopts.debugging_opts.time_passes,
     );
 
+    let ctfe_backtrace = Lock::new(match env::var("RUSTC_CTFE_BACKTRACE") {
+        Ok(ref val) if val == "immediate" => CtfeBacktrace::Immediate,
+        Ok(ref val) if val != "0" => CtfeBacktrace::Capture,
+        _ => CtfeBacktrace::Disabled,
+    });
+
     let sess = Session {
         target: target_cfg,
         host,
@@ -1078,6 +1101,7 @@ fn build_session_(
         trait_methods_not_found: Lock::new(Default::default()),
         confused_type_with_std_module: Lock::new(Default::default()),
         system_library_path: OneThread::new(RefCell::new(Default::default())),
+        ctfe_backtrace,
     };
 
     validate_commandline_args_with_session_available(&sess);