1use std::env;
7use std::env::VarError;
8
9use rustc_ast::token::{self, LitKind};
10use rustc_ast::tokenstream::TokenStream;
11use rustc_ast::{AstDeref, ExprKind, GenericArg, Mutability};
12use rustc_expand::base::{DummyResult, ExpandResult, ExtCtxt, MacEager, MacroExpanderResult};
13use rustc_span::{Ident, Span, Symbol, kw, sym};
14use thin_vec::thin_vec;
15
16use crate::errors;
17use crate::util::{expr_to_string, get_exprs_from_tts, get_single_expr_from_tts};
18
19fn lookup_env<'cx>(cx: &'cx ExtCtxt<'_>, var: Symbol) -> Result<Symbol, VarError> {
20 let var = var.as_str();
21 if let Some(value) = cx.sess.opts.logical_env.get(var) {
22 return Ok(Symbol::intern(value));
23 }
24 Ok(Symbol::intern(&env::var(var)?))
27}
28
29pub(crate) fn expand_option_env<'cx>(
30 cx: &'cx mut ExtCtxt<'_>,
31 sp: Span,
32 tts: TokenStream,
33) -> MacroExpanderResult<'cx> {
34 let ExpandResult::Ready(mac_expr) = get_single_expr_from_tts(cx, sp, tts, "option_env!") else {
35 return ExpandResult::Retry(());
36 };
37 let var_expr = match mac_expr {
38 Ok(var_expr) => var_expr,
39 Err(guar) => return ExpandResult::Ready(DummyResult::any(sp, guar)),
40 };
41 let ExpandResult::Ready(mac) =
42 expr_to_string(cx, var_expr.clone(), "argument must be a string literal")
43 else {
44 return ExpandResult::Retry(());
45 };
46 let var = match mac {
47 Ok((var, _)) => var,
48 Err(guar) => return ExpandResult::Ready(DummyResult::any(sp, guar)),
49 };
50
51 let sp = cx.with_def_site_ctxt(sp);
52 let value = lookup_env(cx, var);
53 cx.sess.psess.env_depinfo.borrow_mut().insert((var, value.as_ref().ok().copied()));
54 let e = match value {
55 Err(VarError::NotPresent) => {
56 let lt = cx.lifetime(sp, Ident::new(kw::StaticLifetime, sp));
57 cx.expr_path(cx.path_all(
58 sp,
59 true,
60 cx.std_path(&[sym::option, sym::Option, sym::None]),
61 vec![GenericArg::Type(cx.ty_ref(
62 sp,
63 cx.ty_ident(sp, Ident::new(sym::str, sp)),
64 Some(lt),
65 Mutability::Not,
66 ))],
67 ))
68 }
69 Err(VarError::NotUnicode(_)) => {
70 let ExprKind::Lit(token::Lit {
71 kind: LitKind::Str | LitKind::StrRaw(..), symbol, ..
72 }) = &var_expr.kind
73 else {
74 unreachable!("`expr_to_string` ensures this is a string lit")
75 };
76
77 let guar = cx.dcx().emit_err(errors::EnvNotUnicode { span: sp, var: *symbol });
78 return ExpandResult::Ready(DummyResult::any(sp, guar));
79 }
80 Ok(value) => cx.expr_call_global(
81 sp,
82 cx.std_path(&[sym::option, sym::Option, sym::Some]),
83 thin_vec![cx.expr_str(sp, value)],
84 ),
85 };
86 ExpandResult::Ready(MacEager::expr(e))
87}
88
89pub(crate) fn expand_env<'cx>(
90 cx: &'cx mut ExtCtxt<'_>,
91 sp: Span,
92 tts: TokenStream,
93) -> MacroExpanderResult<'cx> {
94 let ExpandResult::Ready(mac) = get_exprs_from_tts(cx, tts) else {
95 return ExpandResult::Retry(());
96 };
97 let mut exprs = match mac {
98 Ok(exprs) if exprs.is_empty() || exprs.len() > 2 => {
99 let guar = cx.dcx().emit_err(errors::EnvTakesArgs { span: sp });
100 return ExpandResult::Ready(DummyResult::any(sp, guar));
101 }
102 Err(guar) => return ExpandResult::Ready(DummyResult::any(sp, guar)),
103 Ok(exprs) => exprs.into_iter(),
104 };
105
106 let var_expr = exprs.next().unwrap();
107 let ExpandResult::Ready(mac) = expr_to_string(cx, var_expr.clone(), "expected string literal")
108 else {
109 return ExpandResult::Retry(());
110 };
111 let var = match mac {
112 Ok((var, _)) => var,
113 Err(guar) => return ExpandResult::Ready(DummyResult::any(sp, guar)),
114 };
115
116 let custom_msg = match exprs.next() {
117 None => None,
118 Some(second) => {
119 let ExpandResult::Ready(mac) = expr_to_string(cx, second, "expected string literal")
120 else {
121 return ExpandResult::Retry(());
122 };
123 match mac {
124 Ok((s, _)) => Some(s),
125 Err(guar) => return ExpandResult::Ready(DummyResult::any(sp, guar)),
126 }
127 }
128 };
129
130 let span = cx.with_def_site_ctxt(sp);
131 let value = lookup_env(cx, var);
132 cx.sess.psess.env_depinfo.borrow_mut().insert((var, value.as_ref().ok().copied()));
133 let e = match value {
134 Err(err) => {
135 let ExprKind::Lit(token::Lit {
136 kind: LitKind::Str | LitKind::StrRaw(..), symbol, ..
137 }) = &var_expr.kind
138 else {
139 unreachable!("`expr_to_string` ensures this is a string lit")
140 };
141
142 let guar = match err {
143 VarError::NotPresent => {
144 if let Some(msg_from_user) = custom_msg {
145 cx.dcx()
146 .emit_err(errors::EnvNotDefinedWithUserMessage { span, msg_from_user })
147 } else if is_cargo_env_var(var.as_str()) {
148 cx.dcx().emit_err(errors::EnvNotDefined::CargoEnvVar {
149 span,
150 var: *symbol,
151 var_expr: var_expr.ast_deref(),
152 })
153 } else {
154 cx.dcx().emit_err(errors::EnvNotDefined::CustomEnvVar {
155 span,
156 var: *symbol,
157 var_expr: var_expr.ast_deref(),
158 })
159 }
160 }
161 VarError::NotUnicode(_) => {
162 cx.dcx().emit_err(errors::EnvNotUnicode { span, var: *symbol })
163 }
164 };
165
166 return ExpandResult::Ready(DummyResult::any(sp, guar));
167 }
168 Ok(value) => cx.expr_str(span, value),
169 };
170 ExpandResult::Ready(MacEager::expr(e))
171}
172
173fn is_cargo_env_var(var: &str) -> bool {
175 var.starts_with("CARGO_")
176 || var.starts_with("DEP_")
177 || matches!(var, "OUT_DIR" | "OPT_LEVEL" | "PROFILE" | "HOST" | "TARGET")
178}