proc_macro: Add API for tracked access to environment variables

8 files changed, 66 insertions, 0 deletions

diff --git a/src/libproc_macro/bridge/client.rs b/src/libproc_macro/bridge/client.rs
index 283aa25b0ea..c00e07388b8 100644
--- a/src/libproc_macro/bridge/client.rs
+++ b/src/libproc_macro/bridge/client.rs
@@ -157,6 +157,7 @@ macro_rules! define_handles {
 }
 define_handles! {
     'owned:
+    FreeFunctions,
     TokenStream,
     TokenStreamBuilder,
     TokenStreamIter,
diff --git a/src/libproc_macro/bridge/mod.rs b/src/libproc_macro/bridge/mod.rs
index bf0d8fcee5b..324be9f4701 100644
--- a/src/libproc_macro/bridge/mod.rs
+++ b/src/libproc_macro/bridge/mod.rs
@@ -52,6 +52,10 @@ use std::thread;
 macro_rules! with_api {
     ($S:ident, $self:ident, $m:ident) => {
         $m! {
+            FreeFunctions {
+                fn drop($self: $S::FreeFunctions);
+                fn track_env_var(var: &str, value: Option<&str>);
+            },
             TokenStream {
                 fn drop($self: $S::TokenStream);
                 fn clone($self: &$S::TokenStream) -> $S::TokenStream;
diff --git a/src/libproc_macro/bridge/server.rs b/src/libproc_macro/bridge/server.rs
index ca18d4459aa..eb39025e4c2 100644
--- a/src/libproc_macro/bridge/server.rs
+++ b/src/libproc_macro/bridge/server.rs
@@ -8,6 +8,8 @@ use super::client::HandleStore;
 /// Declare an associated item of one of the traits below, optionally
 /// adjusting it (i.e., adding bounds to types and default bodies to methods).
 macro_rules! associated_item {
+    (type FreeFunctions) =>
+        (type FreeFunctions: 'static;);
     (type TokenStream) =>
         (type TokenStream: 'static + Clone;);
     (type TokenStreamBuilder) =>
diff --git a/src/libproc_macro/lib.rs b/src/libproc_macro/lib.rs
index f960bdecc57..c050a3c591c 100644
--- a/src/libproc_macro/lib.rs
+++ b/src/libproc_macro/lib.rs
@@ -24,6 +24,7 @@
 #![feature(decl_macro)]
 #![feature(extern_types)]
 #![feature(in_band_lifetimes)]
+#![feature(inner_deref)]
 #![feature(negative_impls)]
 #![feature(optin_builtin_traits)]
 #![feature(restricted_std)]
@@ -1160,3 +1161,23 @@ impl fmt::Debug for Literal {
         self.0.fmt(f)
     }
 }
+
+/// Tracked access to environment variables.
+#[unstable(feature = "proc_macro_tracked_env", issue = "74690")]
+pub mod tracked_env {
+    use std::env::{self, VarError};
+    use std::ffi::OsStr;
+
+    /// Retrieve an environment variable and add it to build dependency info.
+    /// Build system executing the compiler will know that the variable was accessed during
+    /// compilation, and will be able to rerun the build when the value of that variable changes.
+    /// Besides the dependency tracking this function should be equivalent to `env::var` from the
+    /// standard library, except that the argument must be UTF-8.
+    #[unstable(feature = "proc_macro_tracked_env", issue = "74690")]
+    pub fn var<K: AsRef<OsStr> + AsRef<str>>(key: K) -> Result<String, VarError> {
+        let key: &str = key.as_ref();
+        let value = env::var(key);
+        crate::bridge::client::FreeFunctions::track_env_var(key, value.as_deref().ok());
+        value
+    }
+}
diff --git a/src/librustc_expand/proc_macro_server.rs b/src/librustc_expand/proc_macro_server.rs
index 2805b4203f9..881d7b84b70 100644
--- a/src/librustc_expand/proc_macro_server.rs
+++ b/src/librustc_expand/proc_macro_server.rs
@@ -274,6 +274,8 @@ impl ToInternal<rustc_errors::Level> for Level {
     }
 }
 
+pub struct FreeFunctions;
+
 #[derive(Clone)]
 pub struct TokenStreamIter {
     cursor: tokenstream::Cursor,
@@ -379,6 +381,7 @@ impl<'a> Rustc<'a> {
 }
 
 impl server::Types for Rustc<'_> {
+    type FreeFunctions = FreeFunctions;
     type TokenStream = TokenStream;
     type TokenStreamBuilder = tokenstream::TokenStreamBuilder;
     type TokenStreamIter = TokenStreamIter;
@@ -392,6 +395,12 @@ impl server::Types for Rustc<'_> {
     type Span = Span;
 }
 
+impl server::FreeFunctions for Rustc<'_> {
+    fn track_env_var(&mut self, var: &str, value: Option<&str>) {
+        self.sess.env_depinfo.borrow_mut().insert((Symbol::intern(var), value.map(Symbol::intern)));
+    }
+}
+
 impl server::TokenStream for Rustc<'_> {
     fn new(&mut self) -> Self::TokenStream {
         TokenStream::default()
diff --git a/src/test/run-make/env-dep-info/Makefile b/src/test/run-make/env-dep-info/Makefile
index 2be0b4b324b..25d9a31c2d6 100644
--- a/src/test/run-make/env-dep-info/Makefile
+++ b/src/test/run-make/env-dep-info/Makefile
@@ -1,8 +1,19 @@
 -include ../../run-make-fulldeps/tools.mk
 
+# FIXME(eddyb) provide `HOST_RUSTC` and `TARGET_RUSTC`
+# instead of hardcoding them everywhere they're needed.
+ifeq ($(IS_MUSL_HOST),1)
+ADDITIONAL_ARGS := $(RUSTFLAGS)
+endif
+
 all:
 	EXISTING_ENV=1 EXISTING_OPT_ENV=1 $(RUSTC) --emit dep-info main.rs
 	$(CGREP) "# env-dep:EXISTING_ENV=1" < $(TMPDIR)/main.d
 	$(CGREP) "# env-dep:EXISTING_OPT_ENV=1" < $(TMPDIR)/main.d
 	$(CGREP) "# env-dep:NONEXISTENT_OPT_ENV" < $(TMPDIR)/main.d
 	$(CGREP) "# env-dep:ESCAPE\nESCAPE\\" < $(TMPDIR)/main.d
+	# Proc macro
+	$(BARE_RUSTC) $(ADDITIONAL_ARGS) --out-dir $(TMPDIR) macro_def.rs
+	EXISTING_PROC_MACRO_ENV=1 $(RUSTC) --emit dep-info macro_use.rs
+	$(CGREP) "# env-dep:EXISTING_PROC_MACRO_ENV=1" < $(TMPDIR)/macro_use.d
+	$(CGREP) "# env-dep:NONEXISTENT_PROC_MACEO_ENV" < $(TMPDIR)/macro_use.d
diff --git a/src/test/run-make/env-dep-info/macro_def.rs b/src/test/run-make/env-dep-info/macro_def.rs
new file mode 100644
index 00000000000..e328eae4832
--- /dev/null
+++ b/src/test/run-make/env-dep-info/macro_def.rs
@@ -0,0 +1,12 @@
+#![feature(proc_macro_tracked_env)]
+#![crate_type = "proc-macro"]
+
+extern crate proc_macro;
+use proc_macro::*;
+
+#[proc_macro]
+pub fn access_env_vars(_: TokenStream) -> TokenStream {
+    let _ = tracked_env::var("EXISTING_PROC_MACRO_ENV");
+    let _ = tracked_env::var("NONEXISTENT_PROC_MACEO_ENV");
+    TokenStream::new()
+}
diff --git a/src/test/run-make/env-dep-info/macro_use.rs b/src/test/run-make/env-dep-info/macro_use.rs
new file mode 100644
index 00000000000..2f5267471fc
--- /dev/null
+++ b/src/test/run-make/env-dep-info/macro_use.rs
@@ -0,0 +1,6 @@
+#[macro_use]
+extern crate macro_def;
+
+access_env_vars!();
+
+fn main() {}