Implement --env compiler flag

Author: GuillaumeGomez

compiler/rustc_builtin_macros/src/env.rs

@@ -13,6 +13,16 @@ use thin_vec::thin_vec;
 
 use crate::errors;
 
+fn lookup_env<'cx>(cx: &'cx ExtCtxt<'_>, var: Symbol) -> Option<Symbol> {
+    let var = var.as_str();
+    if let Some(value) = cx.sess.opts.logical_env.get(var) {
+        return Some(Symbol::intern(value));
+    }
+    // If the environment variable was not defined with the `--env` option, we try to retrieve it
+    // from rustc's environment.
+    env::var(var).ok().as_deref().map(Symbol::intern)
+}
+
 pub fn expand_option_env<'cx>(
     cx: &'cx mut ExtCtxt<'_>,
     sp: Span,
@@ -23,7 +33,7 @@ pub fn expand_option_env<'cx>(
     };
 
     let sp = cx.with_def_site_ctxt(sp);
-    let value = env::var(var.as_str()).ok().as_deref().map(Symbol::intern);
+    let value = lookup_env(cx, var);
     cx.sess.parse_sess.env_depinfo.borrow_mut().insert((var, value));
     let e = match value {
         None => {
@@ -77,7 +87,7 @@ pub fn expand_env<'cx>(
     };
 
     let span = cx.with_def_site_ctxt(sp);
-    let value = env::var(var.as_str()).ok().as_deref().map(Symbol::intern);
+    let value = lookup_env(cx, var);
     cx.sess.parse_sess.env_depinfo.borrow_mut().insert((var, value));
     let e = match value {
         None => {

compiler/rustc_session/src/config.rs

@@ -1115,6 +1115,7 @@ impl Default for Options {
             pretty: None,
             working_dir: RealFileName::LocalPath(std::env::current_dir().unwrap()),
             color: ColorConfig::Auto,
+            logical_env: FxHashMap::default(),
         }
     }
 }
@@ -1812,6 +1813,7 @@ pub fn rustc_optgroups() -> Vec<RustcOptGroup> {
             "Remap source names in all output (compiler messages and output files)",
             "FROM=TO",
         ),
+        opt::multi("", "env", "Inject an environment variable", "VAR=VALUE"),
     ]);
     opts
 }
@@ -2595,6 +2597,23 @@ fn parse_remap_path_prefix(
     mapping
 }
 
+fn parse_logical_env(
+    handler: &mut EarlyErrorHandler,
+    matches: &getopts::Matches,
+) -> FxHashMap<String, String> {
+    let mut vars = FxHashMap::default();
+
+    for arg in matches.opt_strs("env") {
+        if let Some((name, val)) = arg.split_once('=') {
+            vars.insert(name.to_string(), val.to_string());
+        } else {
+            handler.early_error(format!("`--env`: specify value for variable `{arg}`"));
+        }
+    }
+
+    vars
+}
+
 // JUSTIFICATION: before wrapper fn is available
 #[allow(rustc::bad_opt_access)]
 pub fn build_session_options(
@@ -2824,6 +2843,8 @@ pub fn build_session_options(
         handler.early_error("can't dump dependency graph without `-Z query-dep-graph`");
     }
 
+    let logical_env = parse_logical_env(handler, matches);
+
     // Try to find a directory containing the Rust `src`, for more details see
     // the doc comment on the `real_rust_source_base_dir` field.
     let tmp_buf;
@@ -2904,6 +2925,7 @@ pub fn build_session_options(
         pretty,
         working_dir,
         color,
+        logical_env,
     }
 }
 

compiler/rustc_session/src/options.rs

@@ -3,6 +3,7 @@ use crate::config::*;
 use crate::search_paths::SearchPath;
 use crate::utils::NativeLib;
 use crate::{lint, EarlyErrorHandler};
+use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::profiling::TimePassesFormat;
 use rustc_data_structures::stable_hasher::Hash64;
 use rustc_errors::ColorConfig;
@@ -150,6 +151,9 @@ top_level_options!(
 
         target_triple: TargetTriple [TRACKED],
 
+        /// Effective logical environment used by `env!`/`option_env!` macros
+        logical_env: FxHashMap<String, String> [UNTRACKED],
+
         test: bool [TRACKED],
         error_format: ErrorOutputType [UNTRACKED],
         diagnostic_width: Option<usize> [UNTRACKED],