From 6ce2fefea63abab01b9f4892d7608fabba97a7ce Mon Sep 17 00:00:00 2001 From: Guillaume Gomez Date: Mon, 11 Dec 2023 16:44:18 +0100 Subject: [PATCH 1/3] Add support for `--env` on `tracked_env::var` --- compiler/rustc_expand/src/proc_macro_server.rs | 4 ++++ library/proc_macro/src/bridge/mod.rs | 1 + library/proc_macro/src/lib.rs | 3 ++- 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/compiler/rustc_expand/src/proc_macro_server.rs b/compiler/rustc_expand/src/proc_macro_server.rs index 5308e338d7f..523e0eda40e 100644 --- a/compiler/rustc_expand/src/proc_macro_server.rs +++ b/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 { + self.ecx.sess.opts.logical_env.get(var).cloned() + } + fn track_env_var(&mut self, var: &str, value: Option<&str>) { self.sess() .env_depinfo diff --git a/library/proc_macro/src/bridge/mod.rs b/library/proc_macro/src/bridge/mod.rs index 86ce5d9c6d5..75bf3329786 100644 --- a/library/proc_macro/src/bridge/mod.rs +++ b/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; fn track_env_var(var: &str, value: Option<&str>); fn track_path(path: &str); fn literal_from_str(s: &str) -> Result, ()>; diff --git a/library/proc_macro/src/lib.rs b/library/proc_macro/src/lib.rs index 7ebb6810b3e..6e664a162df 100644 --- a/library/proc_macro/src/lib.rs +++ b/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 + AsRef>(key: K) -> Result { 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 } From 5967ceb48df460602ca92fc16e22d558e120cf75 Mon Sep 17 00:00:00 2001 From: Guillaume Gomez Date: Mon, 11 Dec 2023 16:43:35 +0100 Subject: [PATCH 2/3] Update rust-analyzer to support new `injected_env_var` function --- src/tools/rust-analyzer/crates/proc-macro-srv/src/server.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server.rs index fe18451d384..1d36d0c46c5 100644 --- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server.rs +++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server.rs @@ -54,6 +54,10 @@ impl server::Types for RustAnalyzer { } impl server::FreeFunctions for RustAnalyzer { + fn injected_env_var(&mut self, _var: &str) -> Option { + 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 From 453ff1efd5b7644f62bcd2b594eb83225abb58be Mon Sep 17 00:00:00 2001 From: Guillaume Gomez Date: Mon, 11 Dec 2023 16:43:57 +0100 Subject: [PATCH 3/3] Add tests for `--env` usage with `tracked_env::var` --- tests/ui/proc-macro/auxiliary/env.rs | 28 ++++++++++++++++++++++++++++ tests/ui/proc-macro/env.rs | 17 +++++++++++++++++ 2 files changed, 45 insertions(+) create mode 100644 tests/ui/proc-macro/auxiliary/env.rs create mode 100644 tests/ui/proc-macro/env.rs diff --git a/tests/ui/proc-macro/auxiliary/env.rs b/tests/ui/proc-macro/auxiliary/env.rs new file mode 100644 index 00000000000..58bcb08bf06 --- /dev/null +++ b/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() +} diff --git a/tests/ui/proc-macro/env.rs b/tests/ui/proc-macro/env.rs new file mode 100644 index 00000000000..1b1d1873eb3 --- /dev/null +++ b/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); +}