1- use  crate :: assert:: expr_if_not; 
21use  rustc_ast:: { 
32    attr, 
43    ptr:: P , 
54    token, 
65    tokenstream:: { DelimSpan ,  TokenStream ,  TokenTree } , 
7-     BorrowKind ,  Expr ,  ExprKind ,  ItemKind ,  MacArgs ,  MacCall ,  MacDelimiter ,  Mutability ,   Path , 
8-     PathSegment ,  Stmt ,  StructRest ,  UseTree ,  UseTreeKind ,  DUMMY_NODE_ID , 
6+     BinOpKind ,   BorrowKind ,  Expr ,  ExprKind ,  ItemKind ,  MacArgs ,  MacCall ,  MacDelimiter ,  Mutability , 
7+     Path ,   PathSegment ,  Stmt ,  StructRest ,   UnOp ,  UseTree ,  UseTreeKind ,  DUMMY_NODE_ID , 
98} ; 
109use  rustc_ast_pretty:: pprust; 
1110use  rustc_data_structures:: fx:: FxHashSet ; 
@@ -16,11 +15,19 @@ use rustc_span::{
1615} ; 
1716
1817pub ( super )  struct  Context < ' cx ,  ' a >  { 
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 > , 
1923    // Top-level `let captureN = Capture::new()` statements 
2024    capture_decls :  Vec < Capture > , 
2125    cx :  & ' cx  ExtCtxt < ' a > , 
2226    // Formatting string used for debugging 
2327    fmt_string :  String , 
28+     // If the current expression being visited consumes itself. Used to construct 
29+     // `best_case_captures`. 
30+     is_consumed :  bool , 
2431    // Top-level `let __local_bindN = &expr` statements 
2532    local_bind_decls :  Vec < Stmt > , 
2633    // Used to avoid capturing duplicated paths 
@@ -36,9 +43,11 @@ pub(super) struct Context<'cx, 'a> {
3643impl < ' cx ,  ' a >  Context < ' cx ,  ' a >  { 
3744    pub ( super )  fn  new ( cx :  & ' cx  ExtCtxt < ' a > ,  span :  Span )  -> Self  { 
3845        Self  { 
46+             best_case_captures :  <_ >:: default ( ) , 
3947            capture_decls :  <_ >:: default ( ) , 
4048            cx, 
4149            fmt_string :  <_ >:: default ( ) , 
50+             is_consumed :  true , 
4251            local_bind_decls :  <_ >:: default ( ) , 
4352            paths :  <_ >:: default ( ) , 
4453            span, 
@@ -69,14 +78,22 @@ impl<'cx, 'a> Context<'cx, 'a> {
6978        self . manage_cond_expr ( & mut  cond_expr) ; 
7079        let  initial_imports = self . build_initial_imports ( ) ; 
7180        let  panic = self . build_panic ( & expr_str,  panic_path) ; 
81+         let  cond_expr_with_unlikely = self . build_unlikely ( cond_expr) ; 
82+ 
83+         let  Self  {  best_case_captures,  capture_decls,  cx,  local_bind_decls,  span,  .. }  = self ; 
7284
73-         let  Self  {  capture_decls,  cx,  local_bind_decls,  span,  .. }  = self ; 
85+         let  mut  assert_then_stmts = Vec :: with_capacity ( 2 ) ; 
86+         assert_then_stmts. extend ( best_case_captures) ; 
87+         assert_then_stmts. push ( self . cx . stmt_expr ( panic) ) ; 
88+         let  assert_then = self . cx . block ( span,  assert_then_stmts) ; 
7489
7590        let  mut  stmts = Vec :: with_capacity ( 4 ) ; 
7691        stmts. push ( initial_imports) ; 
7792        stmts. extend ( capture_decls. into_iter ( ) . map ( |c| c. decl ) ) ; 
7893        stmts. extend ( local_bind_decls) ; 
79-         stmts. push ( cx. stmt_expr ( expr_if_not ( cx,  span,  cond_expr,  panic,  None ) ) ) ; 
94+         stmts. push ( 
95+             cx. stmt_expr ( cx. expr ( span,  ExprKind :: If ( cond_expr_with_unlikely,  assert_then,  None ) ) ) , 
96+         ) ; 
8097        cx. expr_block ( cx. block ( span,  stmts) ) 
8198    } 
8299
@@ -115,6 +132,16 @@ impl<'cx, 'a> Context<'cx, 'a> {
115132        ) 
116133    } 
117134
135+     /// Takes the conditional expression of `assert!` and then wraps it inside `unlikely` 
136+      fn  build_unlikely ( & self ,  cond_expr :  P < Expr > )  -> P < Expr >  { 
137+         let  unlikely_path = self . cx . std_path ( & [ sym:: intrinsics,  sym:: unlikely] ) ; 
138+         self . cx . expr_call ( 
139+             self . span , 
140+             self . cx . expr_path ( self . cx . path ( self . span ,  unlikely_path) ) , 
141+             vec ! [ self . cx. expr( self . span,  ExprKind :: Unary ( UnOp :: Not ,  cond_expr) ) ] , 
142+         ) 
143+     } 
144+ 
118145    /// The necessary custom `panic!(...)` expression. 
119146     /// 
120147     /// panic!( 
@@ -167,17 +194,39 @@ impl<'cx, 'a> Context<'cx, 'a> {
167194     /// See [Self::manage_initial_capture] and [Self::manage_try_capture] 
168195     fn  manage_cond_expr ( & mut  self ,  expr :  & mut  P < Expr > )  { 
169196        match  ( * expr) . kind  { 
170-             ExprKind :: AddrOf ( _,  _,  ref  mut  local_expr)  => { 
171-                 self . manage_cond_expr ( local_expr) ; 
197+             ExprKind :: AddrOf ( _,  mutability,  ref  mut  local_expr)  => { 
198+                 self . with_is_consumed_management ( 
199+                     matches ! ( mutability,  Mutability :: Mut ) , 
200+                     |this| this. manage_cond_expr ( local_expr) 
201+                 ) ; 
172202            } 
173203            ExprKind :: Array ( ref  mut  local_exprs)  => { 
174204                for  local_expr in  local_exprs { 
175205                    self . manage_cond_expr ( local_expr) ; 
176206                } 
177207            } 
178-             ExprKind :: Binary ( _,  ref  mut  lhs,  ref  mut  rhs)  => { 
179-                 self . manage_cond_expr ( lhs) ; 
180-                 self . manage_cond_expr ( rhs) ; 
208+             ExprKind :: Binary ( ref  op,  ref  mut  lhs,  ref  mut  rhs)  => { 
209+                 self . with_is_consumed_management ( 
210+                     matches ! ( 
211+                         op. node, 
212+                         BinOpKind :: Add 
213+                             | BinOpKind :: And 
214+                             | BinOpKind :: BitAnd 
215+                             | BinOpKind :: BitOr 
216+                             | BinOpKind :: BitXor 
217+                             | BinOpKind :: Div 
218+                             | BinOpKind :: Mul 
219+                             | BinOpKind :: Or 
220+                             | BinOpKind :: Rem 
221+                             | BinOpKind :: Shl 
222+                             | BinOpKind :: Shr 
223+                             | BinOpKind :: Sub 
224+                     ) , 
225+                     |this| { 
226+                         this. manage_cond_expr ( lhs) ; 
227+                         this. manage_cond_expr ( rhs) ; 
228+                     } 
229+                 ) ; 
181230            } 
182231            ExprKind :: Call ( _,  ref  mut  local_exprs)  => { 
183232                for  local_expr in  local_exprs { 
@@ -228,8 +277,11 @@ impl<'cx, 'a> Context<'cx, 'a> {
228277                    self . manage_cond_expr ( local_expr) ; 
229278                } 
230279            } 
231-             ExprKind :: Unary ( _,  ref  mut  local_expr)  => { 
232-                 self . manage_cond_expr ( local_expr) ; 
280+             ExprKind :: Unary ( un_op,  ref  mut  local_expr)  => { 
281+                 self . with_is_consumed_management ( 
282+                     matches ! ( un_op,  UnOp :: Neg  | UnOp :: Not ) , 
283+                     |this| this. manage_cond_expr ( local_expr) 
284+                 ) ; 
233285            } 
234286            // Expressions that are not worth or can not be captured. 
235287            // 
@@ -337,9 +389,23 @@ impl<'cx, 'a> Context<'cx, 'a> {
337389            ) ) 
338390            . add_trailing_semicolon ( ) ; 
339391        let  local_bind_path = self . cx . expr_path ( Path :: from_ident ( local_bind) ) ; 
340-         let  ret = self . cx . stmt_expr ( local_bind_path) ; 
341-         let  block = self . cx . expr_block ( self . cx . block ( self . span ,  vec ! [ try_capture_call,  ret] ) ) ; 
342-         * expr = self . cx . expr_deref ( self . span ,  block) ; 
392+         let  rslt = if  self . is_consumed  { 
393+             let  ret = self . cx . stmt_expr ( local_bind_path) ; 
394+             self . cx . expr_block ( self . cx . block ( self . span ,  vec ! [ try_capture_call,  ret] ) ) 
395+         }  else  { 
396+             self . best_case_captures . push ( try_capture_call) ; 
397+             local_bind_path
398+         } ; 
399+         * expr = self . cx . expr_deref ( self . span ,  rslt) ; 
400+     } 
401+ 
402+     // Calls `f` with the internal `is_consumed` set to `curr_is_consumed` and then 
403+     // sets the internal `is_consumed` back to its original value. 
404+     fn  with_is_consumed_management ( & mut  self ,  curr_is_consumed :  bool ,  f :  impl  FnOnce ( & mut  Self ) )  { 
405+         let  prev_is_consumed = self . is_consumed ; 
406+         self . is_consumed  = curr_is_consumed; 
407+         f ( self ) ; 
408+         self . is_consumed  = prev_is_consumed; 
343409    } 
344410} 
345411
0 commit comments