Auto merge of #118830 - GuillaumeGomez:env-tracked_env, r=Nilstrieb

Add support for `--env` on `tracked_env::var`

Follow-up of https://github.com/rust-lang/rust/pull/118368.
Part of Part of https://github.com/rust-lang/rust/issues/80792.

It adds support of the `--env` option for proc-macros through `tracked_env::var`.

r? `@Nilstrieb`

compiler/rustc_expand/src/proc_macro_server.rs

@@ -426,6 +426,10 @@ impl server::Types for Rustc<'_, '_> {
 }
 
 impl server::FreeFunctions for Rustc<'_, '_> {
+    fn injected_env_var(&mut self, var: &str) -> Option<String> {
+        self.ecx.sess.opts.logical_env.get(var).cloned()
+    }
+
     fn track_env_var(&mut self, var: &str, value: Option<&str>) {
         self.sess()
             .env_depinfo

library/proc_macro/src/bridge/mod.rs

@@ -55,6 +55,7 @@ macro_rules! with_api {
         $m! {
             FreeFunctions {
                 fn drop($self: $S::FreeFunctions);
+                fn injected_env_var(var: &str) -> Option<String>;
                 fn track_env_var(var: &str, value: Option<&str>);
                 fn track_path(path: &str);
                 fn literal_from_str(s: &str) -> Result<Literal<$S::Span, $S::Symbol>, ()>;

library/proc_macro/src/lib.rs

@@ -1503,7 +1503,8 @@ pub mod tracked_env {
     #[unstable(feature = "proc_macro_tracked_env", issue = "99515")]
     pub fn var<K: AsRef<OsStr> + AsRef<str>>(key: K) -> Result<String, VarError> {
         let key: &str = key.as_ref();
-        let value = env::var(key);
+        let value = crate::bridge::client::FreeFunctions::injected_env_var(key)
+            .map_or_else(|| env::var(key), Ok);
         crate::bridge::client::FreeFunctions::track_env_var(key, value.as_deref().ok());
         value
     }

src/tools/rust-analyzer/crates/proc-macro-srv/src/server.rs

@@ -58,6 +58,10 @@ impl server::Types for RustAnalyzer {
 }
 
 impl server::FreeFunctions for RustAnalyzer {
+    fn injected_env_var(&mut self, _var: &str) -> Option<String> {
+        None
+    }
+
     fn track_env_var(&mut self, _var: &str, _value: Option<&str>) {
         // FIXME: track env var accesses
         // https://github.com/rust-lang/rust/pull/71858

tests/ui/proc-macro/auxiliary/env.rs

@@ -0,0 +1,28 @@
+// force-host
+// no-prefer-dynamic
+
+#![crate_type = "proc-macro"]
+#![feature(proc_macro_tracked_env)]
+
+extern crate proc_macro;
+
+use proc_macro::TokenStream;
+use proc_macro::tracked_env::var;
+
+#[proc_macro]
+pub fn generate_const(input: TokenStream) -> TokenStream {
+    let the_const = match var("THE_CONST") {
+        Ok(x) if x == "12" => {
+            "const THE_CONST: u32 = 12;"
+        }
+        _ => {
+            "const THE_CONST: u32 = 0;"
+        }
+    };
+    let another = if var("ANOTHER").is_ok() {
+        "const ANOTHER: u32 = 1;"
+    } else {
+        "const ANOTHER: u32 = 2;"
+    };
+    format!("{the_const}{another}").parse().unwrap()
+}

tests/ui/proc-macro/env.rs

@@ -0,0 +1,17 @@
+// aux-build:env.rs
+// run-pass
+// rustc-env: THE_CONST=1
+// compile-flags: -Zunstable-options --env THE_CONST=12 --env ANOTHER=4
+
+#![crate_name = "foo"]
+
+extern crate env;
+
+use env::generate_const;
+
+generate_const!();
+
+fn main() {
+    assert_eq!(THE_CONST, 12);
+    assert_eq!(ANOTHER, 1);
+}