1+ use  std:: collections:: BTreeMap ; 
2+ use  std:: fmt; 
13use  Context :: * ; 
24
35use  rustc_hir as  hir; 
@@ -25,22 +27,55 @@ enum Context {
2527    Closure ( Span ) , 
2628    Coroutine  {  coroutine_span :  Span ,  kind :  hir:: CoroutineDesugaring ,  source :  hir:: CoroutineSource  } , 
2729    UnlabeledBlock ( Span ) , 
30+     UnlabeledIfBlock ( Span ) , 
2831    LabeledBlock , 
2932    Constant , 
3033} 
3134
32- #[ derive( Copy ,  Clone ) ]  
35+ #[ derive( Clone ) ]  
36+ struct  BlockInfo  { 
37+     name :  String , 
38+     spans :  Vec < Span > , 
39+     suggs :  Vec < Span > , 
40+ } 
41+ 
42+ #[ derive( PartialEq ) ]  
43+ enum  BreakContextKind  { 
44+     Break , 
45+     Continue , 
46+ } 
47+ 
48+ impl  fmt:: Display  for  BreakContextKind  { 
49+     fn  fmt ( & self ,  f :  & mut  fmt:: Formatter < ' _ > )  -> fmt:: Result  { 
50+         match  self  { 
51+             BreakContextKind :: Break  => "break" , 
52+             BreakContextKind :: Continue  => "continue" , 
53+         } 
54+         . fmt ( f) 
55+     } 
56+ } 
57+ 
58+ #[ derive( Clone ) ]  
3359struct  CheckLoopVisitor < ' a ,  ' tcx >  { 
3460    sess :  & ' a  Session , 
3561    tcx :  TyCtxt < ' tcx > , 
36-     cx :  Context , 
62+     // Keep track of a stack of contexts, so that suggestions 
63+     // are not made for contexts where it would be incorrect, 
64+     // such as adding a label for an `if`. 
65+     // e.g. `if 'foo: {}` would be incorrect. 
66+     cx_stack :  Vec < Context > , 
67+     block_breaks :  BTreeMap < Span ,  BlockInfo > , 
3768} 
3869
3970fn  check_mod_loops ( tcx :  TyCtxt < ' _ > ,  module_def_id :  LocalModDefId )  { 
40-     tcx. hir ( ) . visit_item_likes_in_module ( 
41-         module_def_id, 
42-         & mut  CheckLoopVisitor  {  sess :  tcx. sess ,  tcx,  cx :  Normal  } , 
43-     ) ; 
71+     let  mut  check = CheckLoopVisitor  { 
72+         sess :  tcx. sess , 
73+         tcx, 
74+         cx_stack :  vec ! [ Normal ] , 
75+         block_breaks :  Default :: default ( ) , 
76+     } ; 
77+     tcx. hir ( ) . visit_item_likes_in_module ( module_def_id,  & mut  check) ; 
78+     check. report_outside_loop_error ( ) ; 
4479} 
4580
4681pub ( crate )  fn  provide ( providers :  & mut  Providers )  { 
@@ -83,6 +118,45 @@ impl<'a, 'hir> Visitor<'hir> for CheckLoopVisitor<'a, 'hir> {
83118
84119    fn  visit_expr ( & mut  self ,  e :  & ' hir  hir:: Expr < ' hir > )  { 
85120        match  e. kind  { 
121+             hir:: ExprKind :: If ( cond,  then,  else_opt)  => { 
122+                 self . visit_expr ( cond) ; 
123+ 
124+                 let  get_block = |ck_loop :  & CheckLoopVisitor < ' a ,  ' hir > , 
125+                                  expr :  & hir:: Expr < ' hir > |
126+                  -> Option < & hir:: Block < ' hir > >  { 
127+                     if  let  hir:: ExprKind :: Block ( b,  None )  = expr. kind 
128+                         && matches ! ( 
129+                             ck_loop. cx_stack. last( ) , 
130+                             Some ( & Normal ) 
131+                                 | Some ( & Constant ) 
132+                                 | Some ( & UnlabeledBlock ( _) ) 
133+                                 | Some ( & UnlabeledIfBlock ( _) ) 
134+                         ) 
135+                     { 
136+                         Some ( b) 
137+                     }  else  { 
138+                         None 
139+                     } 
140+                 } ; 
141+ 
142+                 if  let  Some ( b)  = get_block ( self ,  then)  { 
143+                     self . with_context ( UnlabeledIfBlock ( b. span . shrink_to_lo ( ) ) ,  |v| { 
144+                         v. visit_block ( b) 
145+                     } ) ; 
146+                 }  else  { 
147+                     self . visit_expr ( then) ; 
148+                 } 
149+ 
150+                 if  let  Some ( else_expr)  = else_opt { 
151+                     if  let  Some ( b)  = get_block ( self ,  else_expr)  { 
152+                         self . with_context ( UnlabeledIfBlock ( b. span . shrink_to_lo ( ) ) ,  |v| { 
153+                             v. visit_block ( b) 
154+                         } ) ; 
155+                     }  else  { 
156+                         self . visit_expr ( else_expr) ; 
157+                     } 
158+                 } 
159+             } 
86160            hir:: ExprKind :: Loop ( ref  b,  _,  source,  _)  => { 
87161                self . with_context ( Loop ( source) ,  |v| v. visit_block ( b) ) ; 
88162            } 
@@ -101,11 +175,14 @@ impl<'a, 'hir> Visitor<'hir> for CheckLoopVisitor<'a, 'hir> {
101175            hir:: ExprKind :: Block ( ref  b,  Some ( _label) )  => { 
102176                self . with_context ( LabeledBlock ,  |v| v. visit_block ( b) ) ; 
103177            } 
104-             hir:: ExprKind :: Block ( ref  b,  None )  if  matches ! ( self . cx ,   Fn )  => { 
178+             hir:: ExprKind :: Block ( ref  b,  None )  if  matches ! ( self . cx_stack . last ( ) ,   Some ( & Fn ) )  => { 
105179                self . with_context ( Normal ,  |v| v. visit_block ( b) ) ; 
106180            } 
107181            hir:: ExprKind :: Block ( ref  b,  None ) 
108-                 if  matches ! ( self . cx,  Normal  | Constant  | UnlabeledBlock ( _) )  =>
182+                 if  matches ! ( 
183+                     self . cx_stack. last( ) , 
184+                     Some ( & Normal )  | Some ( & Constant )  | Some ( & UnlabeledBlock ( _) ) 
185+                 )  =>
109186            { 
110187                self . with_context ( UnlabeledBlock ( b. span . shrink_to_lo ( ) ) ,  |v| v. visit_block ( b) ) ; 
111188            } 
@@ -178,7 +255,12 @@ impl<'a, 'hir> Visitor<'hir> for CheckLoopVisitor<'a, 'hir> {
178255                    Some ( label)  => sp_lo. with_hi ( label. ident . span . hi ( ) ) , 
179256                    None  => sp_lo. shrink_to_lo ( ) , 
180257                } ; 
181-                 self . require_break_cx ( "break" ,  e. span ,  label_sp) ; 
258+                 self . require_break_cx ( 
259+                     BreakContextKind :: Break , 
260+                     e. span , 
261+                     label_sp, 
262+                     self . cx_stack . len ( )  - 1 , 
263+                 ) ; 
182264            } 
183265            hir:: ExprKind :: Continue ( destination)  => { 
184266                self . require_label_in_labeled_block ( e. span ,  & destination,  "continue" ) ; 
@@ -200,7 +282,12 @@ impl<'a, 'hir> Visitor<'hir> for CheckLoopVisitor<'a, 'hir> {
200282                    } 
201283                    Err ( _)  => { } 
202284                } 
203-                 self . require_break_cx ( "continue" ,  e. span ,  e. span ) 
285+                 self . require_break_cx ( 
286+                     BreakContextKind :: Continue , 
287+                     e. span , 
288+                     e. span , 
289+                     self . cx_stack . len ( )  - 1 , 
290+                 ) 
204291            } 
205292            _ => intravisit:: walk_expr ( self ,  e) , 
206293        } 
@@ -212,18 +299,26 @@ impl<'a, 'hir> CheckLoopVisitor<'a, 'hir> {
212299    where 
213300        F :  FnOnce ( & mut  CheckLoopVisitor < ' a ,  ' hir > ) , 
214301    { 
215-         let  old_cx = self . cx ; 
216-         self . cx  = cx; 
302+         self . cx_stack . push ( cx) ; 
217303        f ( self ) ; 
218-         self . cx  = old_cx ; 
304+         self . cx_stack . pop ( ) ; 
219305    } 
220306
221-     fn  require_break_cx ( & self ,  name :  & str ,  span :  Span ,  break_span :  Span )  { 
222-         let  is_break = name == "break" ; 
223-         match  self . cx  { 
307+     fn  require_break_cx ( 
308+         & mut  self , 
309+         br_cx_kind :  BreakContextKind , 
310+         span :  Span , 
311+         break_span :  Span , 
312+         cx_pos :  usize , 
313+     )  { 
314+         match  self . cx_stack [ cx_pos]  { 
224315            LabeledBlock  | Loop ( _)  => { } 
225316            Closure ( closure_span)  => { 
226-                 self . sess . dcx ( ) . emit_err ( BreakInsideClosure  {  span,  closure_span,  name } ) ; 
317+                 self . sess . dcx ( ) . emit_err ( BreakInsideClosure  { 
318+                     span, 
319+                     closure_span, 
320+                     name :  & br_cx_kind. to_string ( ) , 
321+                 } ) ; 
227322            } 
228323            Coroutine  {  coroutine_span,  kind,  source }  => { 
229324                let  kind = match  kind { 
@@ -239,17 +334,32 @@ impl<'a, 'hir> CheckLoopVisitor<'a, 'hir> {
239334                self . sess . dcx ( ) . emit_err ( BreakInsideCoroutine  { 
240335                    span, 
241336                    coroutine_span, 
242-                     name, 
337+                     name :   & br_cx_kind . to_string ( ) , 
243338                    kind, 
244339                    source, 
245340                } ) ; 
246341            } 
247-             UnlabeledBlock ( block_span)  if  is_break && block_span. eq_ctxt ( break_span)  => { 
248-                 let  suggestion = Some ( OutsideLoopSuggestion  {  block_span,  break_span } ) ; 
249-                 self . sess . dcx ( ) . emit_err ( OutsideLoop  {  span,  name,  is_break,  suggestion } ) ; 
342+             UnlabeledBlock ( block_span) 
343+                 if  br_cx_kind == BreakContextKind :: Break  && block_span. eq_ctxt ( break_span)  =>
344+             { 
345+                 let  block = self . block_breaks . entry ( block_span) . or_insert_with ( || BlockInfo  { 
346+                     name :  br_cx_kind. to_string ( ) , 
347+                     spans :  vec ! [ ] , 
348+                     suggs :  vec ! [ ] , 
349+                 } ) ; 
350+                 block. spans . push ( span) ; 
351+                 block. suggs . push ( break_span) ; 
352+             } 
353+             UnlabeledIfBlock ( _)  if  br_cx_kind == BreakContextKind :: Break  => { 
354+                 self . require_break_cx ( br_cx_kind,  span,  break_span,  cx_pos - 1 ) ; 
250355            } 
251-             Normal  | Constant  | Fn  | UnlabeledBlock ( _)  => { 
252-                 self . sess . dcx ( ) . emit_err ( OutsideLoop  {  span,  name,  is_break,  suggestion :  None  } ) ; 
356+             Normal  | Constant  | Fn  | UnlabeledBlock ( _)  | UnlabeledIfBlock ( _)  => { 
357+                 self . sess . dcx ( ) . emit_err ( OutsideLoop  { 
358+                     spans :  vec ! [ span] , 
359+                     name :  & br_cx_kind. to_string ( ) , 
360+                     is_break :  br_cx_kind == BreakContextKind :: Break , 
361+                     suggestion :  None , 
362+                 } ) ; 
253363            } 
254364        } 
255365    } 
@@ -261,12 +371,26 @@ impl<'a, 'hir> CheckLoopVisitor<'a, 'hir> {
261371        cf_type :  & str , 
262372    )  -> bool  { 
263373        if  !span. is_desugaring ( DesugaringKind :: QuestionMark ) 
264-             && self . cx  == LabeledBlock 
374+             && self . cx_stack . last ( )  == Some ( & LabeledBlock ) 
265375            && label. label . is_none ( ) 
266376        { 
267377            self . sess . dcx ( ) . emit_err ( UnlabeledInLabeledBlock  {  span,  cf_type } ) ; 
268378            return  true ; 
269379        } 
270380        false 
271381    } 
382+ 
383+     fn  report_outside_loop_error ( & mut  self )  { 
384+         for  ( s,  block)  in  & self . block_breaks  { 
385+             self . sess . dcx ( ) . emit_err ( OutsideLoop  { 
386+                 spans :  block. spans . clone ( ) , 
387+                 name :  & block. name , 
388+                 is_break :  true , 
389+                 suggestion :  Some ( OutsideLoopSuggestion  { 
390+                     block_span :  * s, 
391+                     break_spans :  block. suggs . clone ( ) , 
392+                 } ) , 
393+             } ) ; 
394+         } 
395+     } 
272396} 
0 commit comments