rustc_builtin_macros/
concat_bytes.rs

1use rustc_ast::ptr::P;
2use rustc_ast::tokenstream::TokenStream;
3use rustc_ast::{ExprKind, LitIntType, LitKind, UintTy, token};
4use rustc_expand::base::{DummyResult, ExpandResult, ExtCtxt, MacEager, MacroExpanderResult};
5use rustc_session::errors::report_lit_error;
6use rustc_span::{ErrorGuaranteed, Span};
7
8use crate::errors;
9use crate::util::get_exprs_from_tts;
10
11/// Emits errors for literal expressions that are invalid inside and outside of an array.
12fn invalid_type_err(
13    cx: &ExtCtxt<'_>,
14    token_lit: token::Lit,
15    span: Span,
16    is_nested: bool,
17) -> ErrorGuaranteed {
18    use errors::{
19        ConcatBytesInvalid, ConcatBytesInvalidSuggestion, ConcatBytesNonU8, ConcatBytesOob,
20    };
21    let snippet = cx.sess.source_map().span_to_snippet(span).ok();
22    let dcx = cx.dcx();
23    match LitKind::from_token_lit(token_lit) {
24        Ok(LitKind::CStr(_, _)) => {
25            // Avoid ambiguity in handling of terminal `NUL` by refusing to
26            // concatenate C string literals as bytes.
27            dcx.emit_err(errors::ConcatCStrLit { span })
28        }
29        Ok(LitKind::Char(_)) => {
30            let sugg =
31                snippet.map(|snippet| ConcatBytesInvalidSuggestion::CharLit { span, snippet });
32            dcx.emit_err(ConcatBytesInvalid { span, lit_kind: "character", sugg })
33        }
34        Ok(LitKind::Str(_, _)) => {
35            // suggestion would be invalid if we are nested
36            let sugg = if !is_nested {
37                snippet.map(|snippet| ConcatBytesInvalidSuggestion::StrLit { span, snippet })
38            } else {
39                None
40            };
41            dcx.emit_err(ConcatBytesInvalid { span, lit_kind: "string", sugg })
42        }
43        Ok(LitKind::Float(_, _)) => {
44            dcx.emit_err(ConcatBytesInvalid { span, lit_kind: "float", sugg: None })
45        }
46        Ok(LitKind::Bool(_)) => {
47            dcx.emit_err(ConcatBytesInvalid { span, lit_kind: "boolean", sugg: None })
48        }
49        Ok(LitKind::Int(_, _)) if !is_nested => {
50            let sugg =
51                snippet.map(|snippet| ConcatBytesInvalidSuggestion::IntLit { span, snippet });
52            dcx.emit_err(ConcatBytesInvalid { span, lit_kind: "numeric", sugg })
53        }
54        Ok(LitKind::Int(val, LitIntType::Unsuffixed | LitIntType::Unsigned(UintTy::U8))) => {
55            assert!(val.get() > u8::MAX.into()); // must be an error
56            dcx.emit_err(ConcatBytesOob { span })
57        }
58        Ok(LitKind::Int(_, _)) => dcx.emit_err(ConcatBytesNonU8 { span }),
59        Ok(LitKind::ByteStr(..) | LitKind::Byte(_)) => unreachable!(),
60        Ok(LitKind::Err(guar)) => guar,
61        Err(err) => report_lit_error(&cx.sess.psess, err, token_lit, span),
62    }
63}
64
65/// Returns `expr` as a *single* byte literal if applicable.
66///
67/// Otherwise, returns `None`, and either pushes the `expr`'s span to `missing_literals` or
68/// updates `guar` accordingly.
69fn handle_array_element(
70    cx: &ExtCtxt<'_>,
71    guar: &mut Option<ErrorGuaranteed>,
72    missing_literals: &mut Vec<rustc_span::Span>,
73    expr: &P<rustc_ast::Expr>,
74) -> Option<u8> {
75    let dcx = cx.dcx();
76
77    match expr.kind {
78        ExprKind::Lit(token_lit) => {
79            match LitKind::from_token_lit(token_lit) {
80                Ok(LitKind::Int(
81                    val,
82                    LitIntType::Unsuffixed | LitIntType::Unsigned(UintTy::U8),
83                )) if let Ok(val) = u8::try_from(val.get()) => {
84                    return Some(val);
85                }
86                Ok(LitKind::Byte(val)) => return Some(val),
87                Ok(LitKind::ByteStr(..)) => {
88                    guar.get_or_insert_with(|| {
89                        dcx.emit_err(errors::ConcatBytesArray { span: expr.span, bytestr: true })
90                    });
91                }
92                _ => {
93                    guar.get_or_insert_with(|| invalid_type_err(cx, token_lit, expr.span, true));
94                }
95            };
96        }
97        ExprKind::Array(_) | ExprKind::Repeat(_, _) => {
98            guar.get_or_insert_with(|| {
99                dcx.emit_err(errors::ConcatBytesArray { span: expr.span, bytestr: false })
100            });
101        }
102        ExprKind::IncludedBytes(..) => {
103            guar.get_or_insert_with(|| {
104                dcx.emit_err(errors::ConcatBytesArray { span: expr.span, bytestr: false })
105            });
106        }
107        _ => missing_literals.push(expr.span),
108    }
109
110    None
111}
112
113pub(crate) fn expand_concat_bytes(
114    cx: &mut ExtCtxt<'_>,
115    sp: Span,
116    tts: TokenStream,
117) -> MacroExpanderResult<'static> {
118    let ExpandResult::Ready(mac) = get_exprs_from_tts(cx, tts) else {
119        return ExpandResult::Retry(());
120    };
121    let es = match mac {
122        Ok(es) => es,
123        Err(guar) => return ExpandResult::Ready(DummyResult::any(sp, guar)),
124    };
125    let mut accumulator = Vec::new();
126    let mut missing_literals = vec![];
127    let mut guar = None;
128    for e in es {
129        match &e.kind {
130            ExprKind::Array(exprs) => {
131                for expr in exprs {
132                    if let Some(elem) =
133                        handle_array_element(cx, &mut guar, &mut missing_literals, expr)
134                    {
135                        accumulator.push(elem);
136                    }
137                }
138            }
139            ExprKind::Repeat(expr, count) => {
140                if let ExprKind::Lit(token_lit) = count.value.kind
141                    && let Ok(LitKind::Int(count_val, _)) = LitKind::from_token_lit(token_lit)
142                {
143                    if let Some(elem) =
144                        handle_array_element(cx, &mut guar, &mut missing_literals, expr)
145                    {
146                        for _ in 0..count_val.get() {
147                            accumulator.push(elem);
148                        }
149                    }
150                } else {
151                    guar = Some(
152                        cx.dcx().emit_err(errors::ConcatBytesBadRepeat { span: count.value.span }),
153                    );
154                }
155            }
156            &ExprKind::Lit(token_lit) => match LitKind::from_token_lit(token_lit) {
157                Ok(LitKind::Byte(val)) => {
158                    accumulator.push(val);
159                }
160                Ok(LitKind::ByteStr(ref bytes, _)) => {
161                    accumulator.extend_from_slice(bytes);
162                }
163                _ => {
164                    guar.get_or_insert_with(|| invalid_type_err(cx, token_lit, e.span, false));
165                }
166            },
167            ExprKind::IncludedBytes(bytes) => {
168                accumulator.extend_from_slice(bytes);
169            }
170            ExprKind::Err(guarantee) => {
171                guar = Some(*guarantee);
172            }
173            ExprKind::Dummy => cx.dcx().span_bug(e.span, "concatenating `ExprKind::Dummy`"),
174            _ => {
175                missing_literals.push(e.span);
176            }
177        }
178    }
179    ExpandResult::Ready(if !missing_literals.is_empty() {
180        let guar = cx.dcx().emit_err(errors::ConcatBytesMissingLiteral { spans: missing_literals });
181        MacEager::expr(DummyResult::raw_expr(sp, Some(guar)))
182    } else if let Some(guar) = guar {
183        MacEager::expr(DummyResult::raw_expr(sp, Some(guar)))
184    } else {
185        let sp = cx.with_def_site_ctxt(sp);
186        MacEager::expr(cx.expr_byte_str(sp, accumulator))
187    })
188}