Add `TyCtx::env_var_os`

Along with `TyCtx::env_var` helper. These can be used to track
environment variable accesses in the query system.

Since `TyCtx::env_var_os` uses `OsStr`, this commit also adds the
necessary trait implementations for that to work.

6 files changed, 75 insertions, 1 deletions

diff --git a/compiler/rustc_data_structures/src/stable_hasher.rs b/compiler/rustc_data_structures/src/stable_hasher.rs
index ffbe54d6206..3a64c924cc2 100644
--- a/compiler/rustc_data_structures/src/stable_hasher.rs
+++ b/compiler/rustc_data_structures/src/stable_hasher.rs
@@ -564,6 +564,8 @@ where
     }
 }
 
+impl_stable_traits_for_trivial_type!(::std::ffi::OsStr);
+
 impl_stable_traits_for_trivial_type!(::std::path::Path);
 impl_stable_traits_for_trivial_type!(::std::path::PathBuf);
 
diff --git a/compiler/rustc_interface/src/passes.rs b/compiler/rustc_interface/src/passes.rs
index 8be7ba7455e..2440f0639c8 100644
--- a/compiler/rustc_interface/src/passes.rs
+++ b/compiler/rustc_interface/src/passes.rs
@@ -1,5 +1,5 @@
 use std::any::Any;
-use std::ffi::OsString;
+use std::ffi::{OsStr, OsString};
 use std::io::{self, BufWriter, Write};
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, LazyLock, OnceLock};
@@ -361,6 +361,31 @@ fn early_lint_checks(tcx: TyCtxt<'_>, (): ()) {
     )
 }
 
+fn env_var_os<'tcx>(tcx: TyCtxt<'tcx>, key: &'tcx OsStr) -> Option<&'tcx OsStr> {
+    let value = env::var_os(key);
+
+    let value_tcx = value.as_ref().map(|value| {
+        let encoded_bytes = tcx.arena.alloc_slice(value.as_encoded_bytes());
+        debug_assert_eq!(value.as_encoded_bytes(), encoded_bytes);
+        // SAFETY: The bytes came from `as_encoded_bytes`, and we assume that
+        // `alloc_slice` is implemented correctly, and passes the same bytes
+        // back (debug asserted above).
+        unsafe { OsStr::from_encoded_bytes_unchecked(encoded_bytes) }
+    });
+
+    // Also add the variable to Cargo's dependency tracking
+    //
+    // NOTE: This only works for passes run before `write_dep_info`. See that
+    // for extension points for configuring environment variables to be
+    // properly change-tracked.
+    tcx.sess.psess.env_depinfo.borrow_mut().insert((
+        Symbol::intern(&key.to_string_lossy()),
+        value.as_ref().and_then(|value| value.to_str()).map(|value| Symbol::intern(&value)),
+    ));
+
+    value_tcx
+}
+
 // Returns all the paths that correspond to generated files.
 fn generated_output_paths(
     tcx: TyCtxt<'_>,
@@ -725,6 +750,7 @@ pub static DEFAULT_QUERY_PROVIDERS: LazyLock<Providers> = LazyLock::new(|| {
         |tcx, _| tcx.arena.alloc_from_iter(tcx.resolutions(()).stripped_cfg_items.steal());
     providers.resolutions = |tcx, ()| tcx.resolver_for_lowering_raw(()).1;
     providers.early_lint_checks = early_lint_checks;
+    providers.env_var_os = env_var_os;
     limits::provide(providers);
     proc_macro_decls::provide(providers);
     rustc_const_eval::provide(providers);
diff --git a/compiler/rustc_middle/src/query/erase.rs b/compiler/rustc_middle/src/query/erase.rs
index 7bbaa0496d5..6c6b9a5510c 100644
--- a/compiler/rustc_middle/src/query/erase.rs
+++ b/compiler/rustc_middle/src/query/erase.rs
@@ -1,3 +1,4 @@
+use std::ffi::OsStr;
 use std::intrinsics::transmute_unchecked;
 use std::mem::MaybeUninit;
 
@@ -67,6 +68,10 @@ impl<T> EraseType for &'_ [T] {
     type Result = [u8; size_of::<&'static [()]>()];
 }
 
+impl EraseType for &'_ OsStr {
+    type Result = [u8; size_of::<&'static OsStr>()];
+}
+
 impl<T> EraseType for &'_ ty::List<T> {
     type Result = [u8; size_of::<&'static ty::List<()>>()];
 }
@@ -174,6 +179,10 @@ impl<T> EraseType for Option<&'_ [T]> {
     type Result = [u8; size_of::<Option<&'static [()]>>()];
 }
 
+impl EraseType for Option<&'_ OsStr> {
+    type Result = [u8; size_of::<Option<&'static OsStr>>()];
+}
+
 impl EraseType for Option<mir::DestructuredConstant<'_>> {
     type Result = [u8; size_of::<Option<mir::DestructuredConstant<'static>>>()];
 }
diff --git a/compiler/rustc_middle/src/query/keys.rs b/compiler/rustc_middle/src/query/keys.rs
index 98314b5abfd..c382bcd726f 100644
--- a/compiler/rustc_middle/src/query/keys.rs
+++ b/compiler/rustc_middle/src/query/keys.rs
@@ -1,5 +1,7 @@
 //! Defines the set of legal keys that can be used in queries.
 
+use std::ffi::OsStr;
+
 use rustc_hir::def_id::{CrateNum, DefId, LOCAL_CRATE, LocalDefId, LocalModDefId, ModDefId};
 use rustc_hir::hir_id::{HirId, OwnerId};
 use rustc_query_system::dep_graph::DepNodeIndex;
@@ -498,6 +500,14 @@ impl Key for Option<Symbol> {
     }
 }
 
+impl<'tcx> Key for &'tcx OsStr {
+    type Cache<V> = DefaultCache<Self, V>;
+
+    fn default_span(&self, _tcx: TyCtxt<'_>) -> Span {
+        DUMMY_SP
+    }
+}
+
 /// Canonical query goals correspond to abstract trait operations that
 /// are not tied to any crate in particular.
 impl<'tcx, T: Clone> Key for CanonicalQueryInput<'tcx, T> {
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index 527c18addbe..850de3f0c2a 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -6,6 +6,7 @@
 
 #![allow(unused_parens)]
 
+use std::ffi::OsStr;
 use std::mem;
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -119,6 +120,21 @@ rustc_queries! {
         desc { "perform lints prior to AST lowering" }
     }
 
+    /// Tracked access to environment variables.
+    ///
+    /// Useful for the implementation of `std::env!`, `proc-macro`s change
+    /// detection and other changes in the compiler's behaviour that is easier
+    /// to control with an environment variable than a flag.
+    ///
+    /// NOTE: This currently does not work with dependency info in the
+    /// analysis, codegen and linking passes, place extra code at the top of
+    /// `rustc_interface::passes::write_dep_info` to make that work.
+    query env_var_os(key: &'tcx OsStr) -> Option<&'tcx OsStr> {
+        // Environment variables are global state
+        eval_always
+        desc { "get the value of an environment variable" }
+    }
+
     query resolutions(_: ()) -> &'tcx ty::ResolverGlobalCtxt {
         no_hash
         desc { "getting the resolver outputs" }
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
index f54dd2b0040..f2003762af3 100644
--- a/compiler/rustc_middle/src/ty/context.rs
+++ b/compiler/rustc_middle/src/ty/context.rs
@@ -7,6 +7,8 @@ pub mod tls;
 use std::assert_matches::{assert_matches, debug_assert_matches};
 use std::borrow::Borrow;
 use std::cmp::Ordering;
+use std::env::VarError;
+use std::ffi::OsStr;
 use std::hash::{Hash, Hasher};
 use std::marker::PhantomData;
 use std::ops::{Bound, Deref};
@@ -1882,6 +1884,15 @@ impl<'tcx> TyCtxt<'tcx> {
         }
         None
     }
+
+    /// Helper to get a tracked environment variable via. [`TyCtxt::env_var_os`] and converting to
+    /// UTF-8 like [`std::env::var`].
+    pub fn env_var<K: ?Sized + AsRef<OsStr>>(self, key: &'tcx K) -> Result<&'tcx str, VarError> {
+        match self.env_var_os(key.as_ref()) {
+            Some(value) => value.to_str().ok_or_else(|| VarError::NotUnicode(value.to_os_string())),
+            None => Err(VarError::NotPresent),
+        }
+    }
 }
 
 impl<'tcx> TyCtxtAt<'tcx> {