compiler/rustc_builtin_macros/src/assert/context.rs

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