add --env option

Author: beepster4096

compiler/rustc_builtin_macros/src/env.rs

@@ -21,7 +21,12 @@ 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 injected_value =
+        cx.sess.opts.injected_env_vars.get(var.as_str()).map(|s| Symbol::intern(s));
+    let value =
+        injected_value.or_else(|| env::var(var.as_str()).ok().as_deref().map(Symbol::intern));
+
     cx.sess.parse_sess.env_depinfo.borrow_mut().insert((var, value));
     let e = match value {
         None => {
@@ -78,7 +83,11 @@ pub fn expand_env<'cx>(
     }
 
     let sp = cx.with_def_site_ctxt(sp);
-    let value = env::var(var.as_str()).ok().as_deref().map(Symbol::intern);
+
+    let key = var.as_str();
+    let injected_value = cx.sess.opts.injected_env_vars.get(key).map(|s| Symbol::intern(s));
+    let value = injected_value.or_else(|| env::var(key).ok().as_deref().map(Symbol::intern));
+
     cx.sess.parse_sess.env_depinfo.borrow_mut().insert((var, value));
     let e = match value {
         None => {

compiler/rustc_session/src/config.rs

@@ -744,6 +744,7 @@ impl Default for Options {
             json_future_incompat: false,
             pretty: None,
             working_dir: RealFileName::LocalPath(std::env::current_dir().unwrap()),
+            injected_env_vars: Default::default(),
         }
     }
 }
@@ -1429,6 +1430,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
 }
@@ -2197,6 +2199,23 @@ fn parse_remap_path_prefix(
     mapping
 }
 
+fn parse_injected_env_vars(
+    matches: &getopts::Matches,
+    error_format: ErrorOutputType,
+) -> 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 {
+            early_error(error_format, &format!("`--env`: specify value for variable `{}`", arg));
+        }
+    }
+
+    vars
+}
+
 // JUSTIFICATION: before wrapper fn is available
 #[cfg_attr(not(bootstrap), allow(rustc::bad_opt_access))]
 pub fn build_session_options(matches: &getopts::Matches) -> Options {
@@ -2470,6 +2489,8 @@ pub fn build_session_options(matches: &getopts::Matches) -> Options {
         RealFileName::LocalPath(path)
     };
 
+    let injected_env_vars = parse_injected_env_vars(matches, error_format);
+
     Options {
         assert_incr_state,
         crate_types,
@@ -2506,6 +2527,7 @@ pub fn build_session_options(matches: &getopts::Matches) -> Options {
         json_future_incompat,
         pretty,
         working_dir,
+        injected_env_vars,
     }
 }
 

compiler/rustc_session/src/options.rs

@@ -10,6 +10,8 @@ use rustc_target::spec::{
     RelocModel, RelroLevel, SplitDebuginfo, StackProtector, TargetTriple, TlsModel,
 };
 
+use rustc_data_structures::fx::FxHashMap;
+
 use rustc_feature::UnstableFeatures;
 use rustc_span::edition::Edition;
 use rustc_span::RealFileName;
@@ -210,6 +212,9 @@ top_level_options!(
 
         /// The (potentially remapped) working directory
         working_dir: RealFileName [TRACKED],
+
+        /// Overridden env vars used for `env!` and `option_env!`
+        injected_env_vars: FxHashMap<String, String> [UNTRACKED],
     }
 );
 

src/test/ui/extenv/extenv-injected.rs

@@ -0,0 +1,8 @@
+// run-pass
+// rustc_env: overridden=no
+// compile-flags: -Z unstable-options --env x=y --env overridden=yes
+
+fn main() {
+    assert_eq!(env!("x"), "y");
+    assert_eq!(env!("overridden"), "yes");
+}