]> git.proxmox.com Git - rustc.git/blame - compiler/rustc_builtin_macros/src/assert/context.rs
New upstream version 1.71.1+dfsg1
[rustc.git] / compiler / rustc_builtin_macros / src / assert / context.rs
CommitLineData
923072b8 1use rustc_ast::{
923072b8
FG
2 ptr::P,
3 token,
4 tokenstream::{DelimSpan, TokenStream, TokenTree},
487cf647
FG
5 BinOpKind, BorrowKind, DelimArgs, Expr, ExprKind, ItemKind, MacCall, MacDelimiter, MethodCall,
6 Mutability, Path, PathSegment, Stmt, StructRest, UnOp, UseTree, UseTreeKind, DUMMY_NODE_ID,
923072b8
FG
7};
8use rustc_ast_pretty::pprust;
9use rustc_data_structures::fx::FxHashSet;
10use rustc_expand::base::ExtCtxt;
11use rustc_span::{
12 symbol::{sym, Ident, Symbol},
13 Span,
14};
9ffffee4 15use thin_vec::{thin_vec, ThinVec};
923072b8
FG
16
17pub(super) struct Context<'cx, 'a> {
064997fb
FG
18 // An optimization.
19 //
20 // Elements that aren't consumed (PartialEq, PartialOrd, ...) can be copied **after** the
21 // `assert!` expression fails rather than copied on-the-fly.
22 best_case_captures: Vec<Stmt>,
923072b8
FG
23 // Top-level `let captureN = Capture::new()` statements
24 capture_decls: Vec<Capture>,
25 cx: &'cx ExtCtxt<'a>,
26 // Formatting string used for debugging
27 fmt_string: String,
064997fb
FG
28 // If the current expression being visited consumes itself. Used to construct
29 // `best_case_captures`.
30 is_consumed: bool,
923072b8
FG
31 // Top-level `let __local_bindN = &expr` statements
32 local_bind_decls: Vec<Stmt>,
33 // Used to avoid capturing duplicated paths
34 //
35 // ```rust
36 // let a = 1i32;
37 // assert!(add(a, a) == 3);
38 // ```
39 paths: FxHashSet<Ident>,
40 span: Span,
41}
42
43impl<'cx, 'a> Context<'cx, 'a> {
44 pub(super) fn new(cx: &'cx ExtCtxt<'a>, span: Span) -> Self {
45 Self {
064997fb 46 best_case_captures: <_>::default(),
923072b8
FG
47 capture_decls: <_>::default(),
48 cx,
49 fmt_string: <_>::default(),
064997fb 50 is_consumed: true,
923072b8
FG
51 local_bind_decls: <_>::default(),
52 paths: <_>::default(),
53 span,
54 }
55 }
56
57 /// Builds the whole `assert!` expression. For example, `let elem = 1; assert!(elem == 1);` expands to:
58 ///
59 /// ```rust
2b03887a 60 /// #![feature(generic_assert_internals)]
923072b8
FG
61 /// let elem = 1;
62 /// {
63 /// #[allow(unused_imports)]
64 /// use ::core::asserting::{TryCaptureGeneric, TryCapturePrintable};
65 /// let mut __capture0 = ::core::asserting::Capture::new();
66 /// let __local_bind0 = &elem;
67 /// if !(
68 /// *{
69 /// (&::core::asserting::Wrapper(__local_bind0)).try_capture(&mut __capture0);
70 /// __local_bind0
71 /// } == 1
72 /// ) {
2b03887a 73 /// panic!("Assertion failed: elem == 1\nWith captures:\n elem = {:?}", __capture0)
923072b8
FG
74 /// }
75 /// }
76 /// ```
77 pub(super) fn build(mut self, mut cond_expr: P<Expr>, panic_path: Path) -> P<Expr> {
78 let expr_str = pprust::expr_to_string(&cond_expr);
79 self.manage_cond_expr(&mut cond_expr);
80 let initial_imports = self.build_initial_imports();
81 let panic = self.build_panic(&expr_str, panic_path);
064997fb
FG
82 let cond_expr_with_unlikely = self.build_unlikely(cond_expr);
83
84 let Self { best_case_captures, capture_decls, cx, local_bind_decls, span, .. } = self;
923072b8 85
9ffffee4 86 let mut assert_then_stmts = ThinVec::with_capacity(2);
064997fb
FG
87 assert_then_stmts.extend(best_case_captures);
88 assert_then_stmts.push(self.cx.stmt_expr(panic));
89 let assert_then = self.cx.block(span, assert_then_stmts);
923072b8 90
9ffffee4 91 let mut stmts = ThinVec::with_capacity(4);
923072b8
FG
92 stmts.push(initial_imports);
93 stmts.extend(capture_decls.into_iter().map(|c| c.decl));
94 stmts.extend(local_bind_decls);
064997fb
FG
95 stmts.push(
96 cx.stmt_expr(cx.expr(span, ExprKind::If(cond_expr_with_unlikely, assert_then, None))),
97 );
923072b8
FG
98 cx.expr_block(cx.block(span, stmts))
99 }
100
101 /// Initial **trait** imports
102 ///
103 /// use ::core::asserting::{ ... };
104 fn build_initial_imports(&self) -> Stmt {
105 let nested_tree = |this: &Self, sym| {
106 (
107 UseTree {
108 prefix: this.cx.path(this.span, vec![Ident::with_dummy_span(sym)]),
487cf647 109 kind: UseTreeKind::Simple(None),
923072b8
FG
110 span: this.span,
111 },
112 DUMMY_NODE_ID,
113 )
114 };
115 self.cx.stmt_item(
116 self.span,
117 self.cx.item(
118 self.span,
119 Ident::empty(),
487cf647 120 thin_vec![self.cx.attr_nested_word(sym::allow, sym::unused_imports, self.span)],
923072b8
FG
121 ItemKind::Use(UseTree {
122 prefix: self.cx.path(self.span, self.cx.std_path(&[sym::asserting])),
9ffffee4 123 kind: UseTreeKind::Nested(thin_vec![
923072b8
FG
124 nested_tree(self, sym::TryCaptureGeneric),
125 nested_tree(self, sym::TryCapturePrintable),
126 ]),
127 span: self.span,
128 }),
129 ),
130 )
131 }
132
064997fb
FG
133 /// Takes the conditional expression of `assert!` and then wraps it inside `unlikely`
134 fn build_unlikely(&self, cond_expr: P<Expr>) -> P<Expr> {
135 let unlikely_path = self.cx.std_path(&[sym::intrinsics, sym::unlikely]);
136 self.cx.expr_call(
137 self.span,
138 self.cx.expr_path(self.cx.path(self.span, unlikely_path)),
9ffffee4 139 thin_vec![self.cx.expr(self.span, ExprKind::Unary(UnOp::Not, cond_expr))],
064997fb
FG
140 )
141 }
142
923072b8
FG
143 /// The necessary custom `panic!(...)` expression.
144 ///
145 /// panic!(
146 /// "Assertion failed: ... \n With expansion: ...",
147 /// __capture0,
148 /// ...
149 /// );
150 fn build_panic(&self, expr_str: &str, panic_path: Path) -> P<Expr> {
151 let escaped_expr_str = escape_to_fmt(expr_str);
152 let initial = [
064997fb 153 TokenTree::token_alone(
923072b8
FG
154 token::Literal(token::Lit {
155 kind: token::LitKind::Str,
156 symbol: Symbol::intern(&if self.fmt_string.is_empty() {
157 format!("Assertion failed: {escaped_expr_str}")
158 } else {
159 format!(
160 "Assertion failed: {escaped_expr_str}\nWith captures:\n{}",
161 &self.fmt_string
162 )
163 }),
164 suffix: None,
165 }),
166 self.span,
167 ),
064997fb 168 TokenTree::token_alone(token::Comma, self.span),
923072b8
FG
169 ];
170 let captures = self.capture_decls.iter().flat_map(|cap| {
171 [
064997fb
FG
172 TokenTree::token_alone(token::Ident(cap.ident.name, false), cap.ident.span),
173 TokenTree::token_alone(token::Comma, self.span),
923072b8
FG
174 ]
175 });
176 self.cx.expr(
177 self.span,
f2b60f7d 178 ExprKind::MacCall(P(MacCall {
923072b8 179 path: panic_path,
487cf647
FG
180 args: P(DelimArgs {
181 dspan: DelimSpan::from_single(self.span),
182 delim: MacDelimiter::Parenthesis,
183 tokens: initial.into_iter().chain(captures).collect::<TokenStream>(),
184 }),
f2b60f7d 185 })),
923072b8
FG
186 )
187 }
188
189 /// Recursive function called until `cond_expr` and `fmt_str` are fully modified.
190 ///
191 /// See [Self::manage_initial_capture] and [Self::manage_try_capture]
192 fn manage_cond_expr(&mut self, expr: &mut P<Expr>) {
487cf647
FG
193 match &mut expr.kind {
194 ExprKind::AddrOf(_, mutability, local_expr) => {
064997fb
FG
195 self.with_is_consumed_management(
196 matches!(mutability, Mutability::Mut),
197 |this| this.manage_cond_expr(local_expr)
198 );
923072b8 199 }
487cf647 200 ExprKind::Array(local_exprs) => {
923072b8
FG
201 for local_expr in local_exprs {
202 self.manage_cond_expr(local_expr);
203 }
204 }
487cf647 205 ExprKind::Binary(op, lhs, rhs) => {
064997fb
FG
206 self.with_is_consumed_management(
207 matches!(
208 op.node,
209 BinOpKind::Add
210 | BinOpKind::And
211 | BinOpKind::BitAnd
212 | BinOpKind::BitOr
213 | BinOpKind::BitXor
214 | BinOpKind::Div
215 | BinOpKind::Mul
216 | BinOpKind::Or
217 | BinOpKind::Rem
218 | BinOpKind::Shl
219 | BinOpKind::Shr
220 | BinOpKind::Sub
221 ),
222 |this| {
223 this.manage_cond_expr(lhs);
224 this.manage_cond_expr(rhs);
225 }
226 );
923072b8 227 }
487cf647 228 ExprKind::Call(_, local_exprs) => {
923072b8
FG
229 for local_expr in local_exprs {
230 self.manage_cond_expr(local_expr);
231 }
232 }
487cf647 233 ExprKind::Cast(local_expr, _) => {
923072b8
FG
234 self.manage_cond_expr(local_expr);
235 }
49aad941
FG
236 ExprKind::If(local_expr, _, _) => {
237 self.manage_cond_expr(local_expr);
238 }
487cf647 239 ExprKind::Index(prefix, suffix) => {
923072b8
FG
240 self.manage_cond_expr(prefix);
241 self.manage_cond_expr(suffix);
242 }
49aad941
FG
243 ExprKind::Let(_, local_expr, _) => {
244 self.manage_cond_expr(local_expr);
245 }
246 ExprKind::Match(local_expr, _) => {
247 self.manage_cond_expr(local_expr);
248 }
487cf647
FG
249 ExprKind::MethodCall(call) => {
250 for arg in &mut call.args {
251 self.manage_cond_expr(arg);
923072b8
FG
252 }
253 }
487cf647 254 ExprKind::Path(_, Path { segments, .. }) if let [path_segment] = &segments[..] => {
923072b8
FG
255 let path_ident = path_segment.ident;
256 self.manage_initial_capture(expr, path_ident);
257 }
487cf647 258 ExprKind::Paren(local_expr) => {
923072b8
FG
259 self.manage_cond_expr(local_expr);
260 }
487cf647
FG
261 ExprKind::Range(prefix, suffix, _) => {
262 if let Some(elem) = prefix {
923072b8
FG
263 self.manage_cond_expr(elem);
264 }
487cf647 265 if let Some(elem) = suffix {
923072b8
FG
266 self.manage_cond_expr(elem);
267 }
268 }
487cf647 269 ExprKind::Repeat(local_expr, elem) => {
923072b8
FG
270 self.manage_cond_expr(local_expr);
271 self.manage_cond_expr(&mut elem.value);
272 }
487cf647 273 ExprKind::Struct(elem) => {
923072b8
FG
274 for field in &mut elem.fields {
275 self.manage_cond_expr(&mut field.expr);
276 }
487cf647 277 if let StructRest::Base(local_expr) = &mut elem.rest {
923072b8
FG
278 self.manage_cond_expr(local_expr);
279 }
280 }
487cf647 281 ExprKind::Tup(local_exprs) => {
923072b8
FG
282 for local_expr in local_exprs {
283 self.manage_cond_expr(local_expr);
284 }
285 }
487cf647 286 ExprKind::Unary(un_op, local_expr) => {
064997fb
FG
287 self.with_is_consumed_management(
288 matches!(un_op, UnOp::Neg | UnOp::Not),
289 |this| this.manage_cond_expr(local_expr)
290 );
923072b8
FG
291 }
292 // Expressions that are not worth or can not be captured.
293 //
294 // Full list instead of `_` to catch possible future inclusions and to
295 // sync with the `rfc-2011-nicer-assert-messages/all-expr-kinds.rs` test.
296 ExprKind::Assign(_, _, _)
297 | ExprKind::AssignOp(_, _, _)
353b0b11 298 | ExprKind::Async(_, _)
49aad941 299 | ExprKind::Await(_, _)
923072b8 300 | ExprKind::Block(_, _)
923072b8 301 | ExprKind::Break(_, _)
487cf647 302 | ExprKind::Closure(_)
923072b8
FG
303 | ExprKind::ConstBlock(_)
304 | ExprKind::Continue(_)
305 | ExprKind::Err
306 | ExprKind::Field(_, _)
307 | ExprKind::ForLoop(_, _, _, _)
49aad941 308 | ExprKind::FormatArgs(_)
487cf647 309 | ExprKind::IncludedBytes(..)
923072b8 310 | ExprKind::InlineAsm(_)
923072b8 311 | ExprKind::Lit(_)
487cf647 312 | ExprKind::Loop(_, _, _)
923072b8 313 | ExprKind::MacCall(_)
49aad941 314 | ExprKind::OffsetOf(_, _)
923072b8
FG
315 | ExprKind::Path(_, _)
316 | ExprKind::Ret(_)
317 | ExprKind::Try(_)
318 | ExprKind::TryBlock(_)
319 | ExprKind::Type(_, _)
320 | ExprKind::Underscore
321 | ExprKind::While(_, _, _)
322 | ExprKind::Yeet(_)
323 | ExprKind::Yield(_) => {}
324 }
325 }
326
327 /// Pushes the top-level declarations and modifies `expr` to try capturing variables.
328 ///
329 /// `fmt_str`, the formatting string used for debugging, is constructed to show possible
330 /// captured variables.
331 fn manage_initial_capture(&mut self, expr: &mut P<Expr>, path_ident: Ident) {
332 if self.paths.contains(&path_ident) {
333 return;
334 } else {
335 self.fmt_string.push_str(" ");
336 self.fmt_string.push_str(path_ident.as_str());
337 self.fmt_string.push_str(" = {:?}\n");
338 let _ = self.paths.insert(path_ident);
339 }
340 let curr_capture_idx = self.capture_decls.len();
341 let capture_string = format!("__capture{curr_capture_idx}");
342 let ident = Ident::new(Symbol::intern(&capture_string), self.span);
343 let init_std_path = self.cx.std_path(&[sym::asserting, sym::Capture, sym::new]);
344 let init = self.cx.expr_call(
345 self.span,
346 self.cx.expr_path(self.cx.path(self.span, init_std_path)),
9ffffee4 347 ThinVec::new(),
923072b8
FG
348 );
349 let capture = Capture { decl: self.cx.stmt_let(self.span, true, ident, init), ident };
350 self.capture_decls.push(capture);
351 self.manage_try_capture(ident, curr_capture_idx, expr);
352 }
353
354 /// Tries to copy `__local_bindN` into `__captureN`.
355 ///
356 /// *{
357 /// (&Wrapper(__local_bindN)).try_capture(&mut __captureN);
358 /// __local_bindN
359 /// }
360 fn manage_try_capture(&mut self, capture: Ident, curr_capture_idx: usize, expr: &mut P<Expr>) {
361 let local_bind_string = format!("__local_bind{curr_capture_idx}");
362 let local_bind = Ident::new(Symbol::intern(&local_bind_string), self.span);
363 self.local_bind_decls.push(self.cx.stmt_let(
364 self.span,
365 false,
366 local_bind,
367 self.cx.expr_addr_of(self.span, expr.clone()),
368 ));
369 let wrapper = self.cx.expr_call(
370 self.span,
371 self.cx.expr_path(
372 self.cx.path(self.span, self.cx.std_path(&[sym::asserting, sym::Wrapper])),
373 ),
9ffffee4 374 thin_vec![self.cx.expr_path(Path::from_ident(local_bind))],
923072b8
FG
375 );
376 let try_capture_call = self
377 .cx
378 .stmt_expr(expr_method_call(
379 self.cx,
380 PathSegment {
381 args: None,
382 id: DUMMY_NODE_ID,
383 ident: Ident::new(sym::try_capture, self.span),
384 },
2b03887a 385 expr_paren(self.cx, self.span, self.cx.expr_addr_of(self.span, wrapper)),
9ffffee4 386 thin_vec![expr_addr_of_mut(
2b03887a
FG
387 self.cx,
388 self.span,
389 self.cx.expr_path(Path::from_ident(capture)),
390 )],
923072b8
FG
391 self.span,
392 ))
393 .add_trailing_semicolon();
394 let local_bind_path = self.cx.expr_path(Path::from_ident(local_bind));
064997fb
FG
395 let rslt = if self.is_consumed {
396 let ret = self.cx.stmt_expr(local_bind_path);
9ffffee4 397 self.cx.expr_block(self.cx.block(self.span, thin_vec![try_capture_call, ret]))
064997fb
FG
398 } else {
399 self.best_case_captures.push(try_capture_call);
400 local_bind_path
401 };
402 *expr = self.cx.expr_deref(self.span, rslt);
403 }
404
405 // Calls `f` with the internal `is_consumed` set to `curr_is_consumed` and then
406 // sets the internal `is_consumed` back to its original value.
407 fn with_is_consumed_management(&mut self, curr_is_consumed: bool, f: impl FnOnce(&mut Self)) {
408 let prev_is_consumed = self.is_consumed;
409 self.is_consumed = curr_is_consumed;
410 f(self);
411 self.is_consumed = prev_is_consumed;
923072b8
FG
412 }
413}
414
415/// Information about a captured element.
416#[derive(Debug)]
417struct Capture {
418 // Generated indexed `Capture` statement.
419 //
420 // `let __capture{} = Capture::new();`
421 decl: Stmt,
422 // The name of the generated indexed `Capture` variable.
423 //
424 // `__capture{}`
425 ident: Ident,
426}
427
428/// Escapes to use as a formatting string.
429fn escape_to_fmt(s: &str) -> String {
430 let mut rslt = String::with_capacity(s.len());
431 for c in s.chars() {
432 rslt.extend(c.escape_debug());
433 match c {
434 '{' | '}' => rslt.push(c),
435 _ => {}
436 }
437 }
438 rslt
439}
440
441fn expr_addr_of_mut(cx: &ExtCtxt<'_>, sp: Span, e: P<Expr>) -> P<Expr> {
442 cx.expr(sp, ExprKind::AddrOf(BorrowKind::Ref, Mutability::Mut, e))
443}
444
445fn expr_method_call(
446 cx: &ExtCtxt<'_>,
487cf647 447 seg: PathSegment,
2b03887a 448 receiver: P<Expr>,
9ffffee4 449 args: ThinVec<P<Expr>>,
923072b8
FG
450 span: Span,
451) -> P<Expr> {
487cf647 452 cx.expr(span, ExprKind::MethodCall(Box::new(MethodCall { seg, receiver, args, span })))
923072b8
FG
453}
454
455fn expr_paren(cx: &ExtCtxt<'_>, sp: Span, e: P<Expr>) -> P<Expr> {
456 cx.expr(sp, ExprKind::Paren(e))
457}