4 tokenstream::{DelimSpan, TokenStream, TokenTree},
5 BinOpKind, BorrowKind, DelimArgs, Expr, ExprKind, ItemKind, MacCall, MacDelimiter, MethodCall,
6 Mutability, Path, PathSegment, Stmt, StructRest, UnOp, UseTree, UseTreeKind, DUMMY_NODE_ID,
8 use rustc_ast_pretty::pprust;
9 use rustc_data_structures::fx::FxHashSet;
10 use rustc_expand::base::ExtCtxt;
12 symbol::{sym, Ident, Symbol},
15 use thin_vec::thin_vec;
17 pub(super) struct Context<'cx, 'a> {
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>,
23 // Top-level `let captureN = Capture::new()` statements
24 capture_decls: Vec<Capture>,
26 // Formatting string used for debugging
28 // If the current expression being visited consumes itself. Used to construct
29 // `best_case_captures`.
31 // Top-level `let __local_bindN = &expr` statements
32 local_bind_decls: Vec<Stmt>,
33 // Used to avoid capturing duplicated paths
37 // assert!(add(a, a) == 3);
39 paths: FxHashSet<Ident>,
43 impl<'cx, 'a> Context<'cx, 'a> {
44 pub(super) fn new(cx: &'cx ExtCtxt<'a>, span: Span) -> Self {
46 best_case_captures: <_>::default(),
47 capture_decls: <_>::default(),
49 fmt_string: <_>::default(),
51 local_bind_decls: <_>::default(),
52 paths: <_>::default(),
57 /// Builds the whole `assert!` expression. For example, `let elem = 1; assert!(elem == 1);` expands to:
60 /// #![feature(generic_assert_internals)]
63 /// #[allow(unused_imports)]
64 /// use ::core::asserting::{TryCaptureGeneric, TryCapturePrintable};
65 /// let mut __capture0 = ::core::asserting::Capture::new();
66 /// let __local_bind0 = &elem;
69 /// (&::core::asserting::Wrapper(__local_bind0)).try_capture(&mut __capture0);
73 /// panic!("Assertion failed: elem == 1\nWith captures:\n elem = {:?}", __capture0)
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);
82 let cond_expr_with_unlikely = self.build_unlikely(cond_expr);
84 let Self { best_case_captures, capture_decls, cx, local_bind_decls, span, .. } = self;
86 let mut assert_then_stmts = Vec::with_capacity(2);
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);
91 let mut stmts = Vec::with_capacity(4);
92 stmts.push(initial_imports);
93 stmts.extend(capture_decls.into_iter().map(|c| c.decl));
94 stmts.extend(local_bind_decls);
96 cx.stmt_expr(cx.expr(span, ExprKind::If(cond_expr_with_unlikely, assert_then, None))),
98 cx.expr_block(cx.block(span, stmts))
101 /// Initial **trait** imports
103 /// use ::core::asserting::{ ... };
104 fn build_initial_imports(&self) -> Stmt {
105 let nested_tree = |this: &Self, sym| {
108 prefix: this.cx.path(this.span, vec![Ident::with_dummy_span(sym)]),
109 kind: UseTreeKind::Simple(None),
120 thin_vec![self.cx.attr_nested_word(sym::allow, sym::unused_imports, self.span)],
121 ItemKind::Use(UseTree {
122 prefix: self.cx.path(self.span, self.cx.std_path(&[sym::asserting])),
123 kind: UseTreeKind::Nested(vec![
124 nested_tree(self, sym::TryCaptureGeneric),
125 nested_tree(self, sym::TryCapturePrintable),
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]);
138 self.cx.expr_path(self.cx.path(self.span, unlikely_path)),
139 vec![self.cx.expr(self.span, ExprKind::Unary(UnOp::Not, cond_expr))],
143 /// The necessary custom `panic!(...)` expression.
146 /// "Assertion failed: ... \n With expansion: ...",
150 fn build_panic(&self, expr_str: &str, panic_path: Path) -> P<Expr> {
151 let escaped_expr_str = escape_to_fmt(expr_str);
153 TokenTree::token_alone(
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}")
160 "Assertion failed: {escaped_expr_str}\nWith captures:\n{}",
168 TokenTree::token_alone(token::Comma, self.span),
170 let captures = self.capture_decls.iter().flat_map(|cap| {
172 TokenTree::token_alone(token::Ident(cap.ident.name, false), cap.ident.span),
173 TokenTree::token_alone(token::Comma, self.span),
178 ExprKind::MacCall(P(MacCall {
181 dspan: DelimSpan::from_single(self.span),
182 delim: MacDelimiter::Parenthesis,
183 tokens: initial.into_iter().chain(captures).collect::<TokenStream>(),
185 prior_type_ascription: None,
190 /// Recursive function called until `cond_expr` and `fmt_str` are fully modified.
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)
201 ExprKind::Array(local_exprs) => {
202 for local_expr in local_exprs {
203 self.manage_cond_expr(local_expr);
206 ExprKind::Binary(op, lhs, rhs) => {
207 self.with_is_consumed_management(
224 this.manage_cond_expr(lhs);
225 this.manage_cond_expr(rhs);
229 ExprKind::Call(_, local_exprs) => {
230 for local_expr in local_exprs {
231 self.manage_cond_expr(local_expr);
234 ExprKind::Cast(local_expr, _) => {
235 self.manage_cond_expr(local_expr);
237 ExprKind::Index(prefix, suffix) => {
238 self.manage_cond_expr(prefix);
239 self.manage_cond_expr(suffix);
241 ExprKind::MethodCall(call) => {
242 for arg in &mut call.args {
243 self.manage_cond_expr(arg);
246 ExprKind::Path(_, Path { segments, .. }) if let [path_segment] = &segments[..] => {
247 let path_ident = path_segment.ident;
248 self.manage_initial_capture(expr, path_ident);
250 ExprKind::Paren(local_expr) => {
251 self.manage_cond_expr(local_expr);
253 ExprKind::Range(prefix, suffix, _) => {
254 if let Some(elem) = prefix {
255 self.manage_cond_expr(elem);
257 if let Some(elem) = suffix {
258 self.manage_cond_expr(elem);
261 ExprKind::Repeat(local_expr, elem) => {
262 self.manage_cond_expr(local_expr);
263 self.manage_cond_expr(&mut elem.value);
265 ExprKind::Struct(elem) => {
266 for field in &mut elem.fields {
267 self.manage_cond_expr(&mut field.expr);
269 if let StructRest::Base(local_expr) = &mut elem.rest {
270 self.manage_cond_expr(local_expr);
273 ExprKind::Tup(local_exprs) => {
274 for local_expr in local_exprs {
275 self.manage_cond_expr(local_expr);
278 ExprKind::Unary(un_op, local_expr) => {
279 self.with_is_consumed_management(
280 matches!(un_op, UnOp::Neg | UnOp::Not),
281 |this| this.manage_cond_expr(local_expr)
284 // Expressions that are not worth or can not be captured.
286 // Full list instead of `_` to catch possible future inclusions and to
287 // sync with the `rfc-2011-nicer-assert-messages/all-expr-kinds.rs` test.
288 ExprKind::Assign(_, _, _)
289 | ExprKind::AssignOp(_, _, _)
290 | ExprKind::Async(_, _, _)
292 | ExprKind::Block(_, _)
294 | ExprKind::Break(_, _)
295 | ExprKind::Closure(_)
296 | ExprKind::ConstBlock(_)
297 | ExprKind::Continue(_)
299 | ExprKind::Field(_, _)
300 | ExprKind::FormatArgs(_)
301 | ExprKind::ForLoop(_, _, _, _)
302 | ExprKind::If(_, _, _)
303 | ExprKind::IncludedBytes(..)
304 | ExprKind::InlineAsm(_)
305 | ExprKind::Let(_, _, _)
307 | ExprKind::Loop(_, _, _)
308 | ExprKind::MacCall(_)
309 | ExprKind::Match(_, _)
310 | ExprKind::Path(_, _)
313 | ExprKind::TryBlock(_)
314 | ExprKind::Type(_, _)
315 | ExprKind::Underscore
316 | ExprKind::While(_, _, _)
318 | ExprKind::Yield(_) => {}
322 /// Pushes the top-level declarations and modifies `expr` to try capturing variables.
324 /// `fmt_str`, the formatting string used for debugging, is constructed to show possible
325 /// captured variables.
326 fn manage_initial_capture(&mut self, expr: &mut P<Expr>, path_ident: Ident) {
327 if self.paths.contains(&path_ident) {
330 self.fmt_string.push_str(" ");
331 self.fmt_string.push_str(path_ident.as_str());
332 self.fmt_string.push_str(" = {:?}\n");
333 let _ = self.paths.insert(path_ident);
335 let curr_capture_idx = self.capture_decls.len();
336 let capture_string = format!("__capture{curr_capture_idx}");
337 let ident = Ident::new(Symbol::intern(&capture_string), self.span);
338 let init_std_path = self.cx.std_path(&[sym::asserting, sym::Capture, sym::new]);
339 let init = self.cx.expr_call(
341 self.cx.expr_path(self.cx.path(self.span, init_std_path)),
344 let capture = Capture { decl: self.cx.stmt_let(self.span, true, ident, init), ident };
345 self.capture_decls.push(capture);
346 self.manage_try_capture(ident, curr_capture_idx, expr);
349 /// Tries to copy `__local_bindN` into `__captureN`.
352 /// (&Wrapper(__local_bindN)).try_capture(&mut __captureN);
355 fn manage_try_capture(&mut self, capture: Ident, curr_capture_idx: usize, expr: &mut P<Expr>) {
356 let local_bind_string = format!("__local_bind{curr_capture_idx}");
357 let local_bind = Ident::new(Symbol::intern(&local_bind_string), self.span);
358 self.local_bind_decls.push(self.cx.stmt_let(
362 self.cx.expr_addr_of(self.span, expr.clone()),
364 let wrapper = self.cx.expr_call(
367 self.cx.path(self.span, self.cx.std_path(&[sym::asserting, sym::Wrapper])),
369 vec![self.cx.expr_path(Path::from_ident(local_bind))],
371 let try_capture_call = self
373 .stmt_expr(expr_method_call(
378 ident: Ident::new(sym::try_capture, self.span),
380 expr_paren(self.cx, self.span, self.cx.expr_addr_of(self.span, wrapper)),
381 vec![expr_addr_of_mut(
384 self.cx.expr_path(Path::from_ident(capture)),
388 .add_trailing_semicolon();
389 let local_bind_path = self.cx.expr_path(Path::from_ident(local_bind));
390 let rslt = if self.is_consumed {
391 let ret = self.cx.stmt_expr(local_bind_path);
392 self.cx.expr_block(self.cx.block(self.span, vec![try_capture_call, ret]))
394 self.best_case_captures.push(try_capture_call);
397 *expr = self.cx.expr_deref(self.span, rslt);
400 // Calls `f` with the internal `is_consumed` set to `curr_is_consumed` and then
401 // sets the internal `is_consumed` back to its original value.
402 fn with_is_consumed_management(&mut self, curr_is_consumed: bool, f: impl FnOnce(&mut Self)) {
403 let prev_is_consumed = self.is_consumed;
404 self.is_consumed = curr_is_consumed;
406 self.is_consumed = prev_is_consumed;
410 /// Information about a captured element.
413 // Generated indexed `Capture` statement.
415 // `let __capture{} = Capture::new();`
417 // The name of the generated indexed `Capture` variable.
423 /// Escapes to use as a formatting string.
424 fn escape_to_fmt(s: &str) -> String {
425 let mut rslt = String::with_capacity(s.len());
427 rslt.extend(c.escape_debug());
429 '{' | '}' => rslt.push(c),
436 fn expr_addr_of_mut(cx: &ExtCtxt<'_>, sp: Span, e: P<Expr>) -> P<Expr> {
437 cx.expr(sp, ExprKind::AddrOf(BorrowKind::Ref, Mutability::Mut, e))
447 cx.expr(span, ExprKind::MethodCall(Box::new(MethodCall { seg, receiver, args, span })))
450 fn expr_paren(cx: &ExtCtxt<'_>, sp: Span, e: P<Expr>) -> P<Expr> {
451 cx.expr(sp, ExprKind::Paren(e))