Fix bug where option_env! would return None when env var is present but not valid Unicode

Author: beetrees

Diff for: compiler/rustc_builtin_macros/src/env.rs

@@ -6,7 +6,7 @@
 use rustc_ast::token::{self, LitKind};
 use rustc_ast::tokenstream::TokenStream;
 use rustc_ast::{AstDeref, ExprKind, GenericArg, Mutability};
-use rustc_expand::base::{expr_to_string, get_exprs_from_tts, get_single_str_from_tts};
+use rustc_expand::base::{expr_to_string, get_exprs_from_tts, get_single_expr_from_tts};
 use rustc_expand::base::{DummyResult, ExpandResult, ExtCtxt, MacEager, MacroExpanderResult};
 use rustc_span::symbol::{kw, sym, Ident, Symbol};
 use rustc_span::Span;
@@ -31,19 +31,28 @@ pub fn expand_option_env<'cx>(
     sp: Span,
     tts: TokenStream,
 ) -> MacroExpanderResult<'cx> {
-    let ExpandResult::Ready(mac) = get_single_str_from_tts(cx, sp, tts, "option_env!") else {
+    let ExpandResult::Ready(mac_expr) = get_single_expr_from_tts(cx, sp, tts, "option_env!") else {
+        return ExpandResult::Retry(());
+    };
+    let var_expr = match mac_expr {
+        Ok(var_expr) => var_expr,
+        Err(guar) => return ExpandResult::Ready(DummyResult::any(sp, guar)),
+    };
+    let ExpandResult::Ready(mac) =
+        expr_to_string(cx, var_expr.clone(), "argument must be a string literal")
+    else {
         return ExpandResult::Retry(());
     };
     let var = match mac {
-        Ok(var) => var,
+        Ok((var, _)) => var,
         Err(guar) => return ExpandResult::Ready(DummyResult::any(sp, guar)),
     };
 
     let sp = cx.with_def_site_ctxt(sp);
-    let value = lookup_env(cx, var).ok();
-    cx.sess.psess.env_depinfo.borrow_mut().insert((var, value));
+    let value = lookup_env(cx, var);
+    cx.sess.psess.env_depinfo.borrow_mut().insert((var, value.as_ref().ok().copied()));
     let e = match value {
-        None => {
+        Err(VarError::NotPresent) => {
             let lt = cx.lifetime(sp, Ident::new(kw::StaticLifetime, sp));
             cx.expr_path(cx.path_all(
                 sp,
@@ -57,7 +66,18 @@ pub fn expand_option_env<'cx>(
                 ))],
             ))
         }
-        Some(value) => cx.expr_call_global(
+        Err(VarError::NotUnicode(_)) => {
+            let ExprKind::Lit(token::Lit {
+                kind: LitKind::Str | LitKind::StrRaw(..), symbol, ..
+            }) = &var_expr.kind
+            else {
+                unreachable!("`expr_to_string` ensures this is a string lit")
+            };
+
+            let guar = cx.dcx().emit_err(errors::EnvNotUnicode { span: sp, var: *symbol });
+            return ExpandResult::Ready(DummyResult::any(sp, guar));
+        }
+        Ok(value) => cx.expr_call_global(
             sp,
             cx.std_path(&[sym::option, sym::Option, sym::Some]),
             thin_vec![cx.expr_str(sp, value)],

Diff for: compiler/rustc_expand/src/base.rs

@@ -1379,6 +1379,25 @@ pub fn get_single_str_from_tts(
     tts: TokenStream,
     name: &str,
 ) -> ExpandResult<Result<Symbol, ErrorGuaranteed>, ()> {
+    let ExpandResult::Ready(ret) = get_single_expr_from_tts(cx, span, tts, name) else {
+        return ExpandResult::Retry(());
+    };
+    match ret {
+        Ok(ret) => {
+            expr_to_string(cx, ret, "argument must be a string literal").map(|s| s.map(|(s, _)| s))
+        }
+        Err(e) => ExpandResult::Ready(Err(e)),
+    }
+}
+
+/// Interpreting `tts` as a comma-separated sequence of expressions,
+/// expect exactly one expression, or emit an error and return `Err`.
+pub fn get_single_expr_from_tts(
+    cx: &mut ExtCtxt<'_>,
+    span: Span,
+    tts: TokenStream,
+    name: &str,
+) -> ExpandResult<Result<P<ast::Expr>, ErrorGuaranteed>, ()> {
     let mut p = cx.new_parser_from_tts(tts);
     if p.token == token::Eof {
         let guar = cx.dcx().emit_err(errors::OnlyOneArgument { span, name });
@@ -1393,7 +1412,7 @@ pub fn get_single_str_from_tts(
     if p.token != token::Eof {
         cx.dcx().emit_err(errors::OnlyOneArgument { span, name });
     }
-    expr_to_string(cx, ret, "argument must be a string literal").map(|s| s.map(|(s, _)| s))
+    ExpandResult::Ready(Ok(ret))
 }
 
 /// Extracts comma-separated expressions from `tts`.

Diff for: library/core/src/macros/mod.rs

@@ -1089,17 +1089,19 @@ pub(crate) mod builtin {
     ///
     /// If the named environment variable is present at compile time, this will
     /// expand into an expression of type `Option<&'static str>` whose value is
-    /// `Some` of the value of the environment variable. If the environment
-    /// variable is not present, then this will expand to `None`. See
-    /// [`Option<T>`][Option] for more information on this type.  Use
-    /// [`std::env::var`] instead if you want to read the value at runtime.
+    /// `Some` of the value of the environment variable (a compilation error
+    /// will be emitted if the environment variable is not a valid Unicode
+    /// string). If the environment variable is not present, then this will
+    /// expand to `None`. See [`Option<T>`][Option] for more information on this
+    /// type.  Use [`std::env::var`] instead if you want to read the value at
+    /// runtime.
     ///
     /// [`std::env::var`]: ../std/env/fn.var.html
     ///
-    /// A compile time error is never emitted when using this macro regardless
-    /// of whether the environment variable is present or not.
-    /// To emit a compile error if the environment variable is not present,
-    /// use the [`env!`] macro instead.
+    /// A compile time error is only emitted when using this macro if the
+    /// environment variable exists and is not a valid Unicode string. To also
+    /// emit a compile error if the environment variable is not present, use the
+    /// [`env!`] macro instead.
     ///
     /// # Examples
     ///

Diff for: tests/run-make/non-unicode-env/non_unicode_env.rs

@@ -1,3 +1,4 @@
 fn main() {
     let _ = env!("NON_UNICODE_VAR");
+    let _ = option_env!("NON_UNICODE_VAR");
 }

Diff for: tests/run-make/non-unicode-env/non_unicode_env.stderr

@@ -6,5 +6,13 @@ error: environment variable `NON_UNICODE_VAR` is not a valid Unicode string
   |
   = note: this error originates in the macro `env` (in Nightly builds, run with -Z macro-backtrace for more info)
 
-error: aborting due to 1 previous error
+error: environment variable `NON_UNICODE_VAR` is not a valid Unicode string
+ --> non_unicode_env.rs:3:13
+  |
+3 |     let _ = option_env!("NON_UNICODE_VAR");
+  |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  |
+  = note: this error originates in the macro `option_env` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+error: aborting due to 2 previous errors