Properly implement labeled breaks in while conditions

src/librustc/cfg/construct.rs

@@ -220,15 +220,24 @@ impl<'a, 'tcx> CFGBuilder<'a, 'tcx> {
                 // Note that `break` and `continue` statements
                 // may cause additional edges.
 
-                // Is the condition considered part of the loop?
                 let loopback = self.add_dummy_node(&[pred]);              // 1
-                let cond_exit = self.expr(&cond, loopback);             // 2
-                let expr_exit = self.add_ast_node(expr.id, &[cond_exit]); // 3
+
+                // Create expr_exit without pred (cond_exit)
+                let expr_exit = self.add_ast_node(expr.id, &[]);         // 3
+
+                // The LoopScope needs to be on the loop_scopes stack while evaluating the
+                // condition and the body of the loop (both can break out of the loop)
                 self.loop_scopes.push(LoopScope {
                     loop_id: expr.id,
                     continue_index: loopback,
                     break_index: expr_exit
                 });
+
+                let cond_exit = self.expr(&cond, loopback);             // 2
+
+                // Add pred (cond_exit) to expr_exit
+                self.add_contained_edge(cond_exit, expr_exit);
+
                 let body_exit = self.block(&body, cond_exit);          // 4
                 self.add_contained_edge(body_exit, loopback);            // 5
                 self.loop_scopes.pop();
@@ -580,11 +589,17 @@ impl<'a, 'tcx> CFGBuilder<'a, 'tcx> {
     fn find_scope(&self,
                   expr: &hir::Expr,
                   label: hir::Label) -> LoopScope {
-        for l in &self.loop_scopes {
-            if l.loop_id == label.loop_id {
-                return *l;
+
+        match label.loop_id.into() {
+            Ok(loop_id) => {
+                for l in &self.loop_scopes {
+                    if l.loop_id == loop_id {
+                        return *l;
+                    }
+                }
+                span_bug!(expr.span, "no loop scope for id {}", loop_id);
             }
+            Err(err) => span_bug!(expr.span, "loop scope error: {}",  err)
         }
-        span_bug!(expr.span, "no loop scope for id {}", label.loop_id);
     }
 }

src/librustc/hir/intravisit.rs

@@ -1008,14 +1008,18 @@ pub fn walk_expr<'v, V: Visitor<'v>>(visitor: &mut V, expression: &'v Expr) {
         }
         ExprBreak(label, ref opt_expr) => {
             label.ident.map(|ident| {
-                visitor.visit_def_mention(Def::Label(label.loop_id));
+                if let Ok(loop_id) = label.loop_id.into() {
+                    visitor.visit_def_mention(Def::Label(loop_id));
+                }
                 visitor.visit_name(ident.span, ident.node.name);
             });
             walk_list!(visitor, visit_expr, opt_expr);
         }
         ExprAgain(label) => {
             label.ident.map(|ident| {
-                visitor.visit_def_mention(Def::Label(label.loop_id));
+                if let Ok(loop_id) = label.loop_id.into() {
+                    visitor.visit_def_mention(Def::Label(loop_id));
+                }
                 visitor.visit_name(ident.span, ident.node.name);
             });
         }

src/librustc/hir/lowering.rs

@@ -81,6 +81,7 @@ pub struct LoweringContext<'a> {
     bodies: FxHashMap<hir::BodyId, hir::Body>,
 
     loop_scopes: Vec<NodeId>,
+    is_in_loop_condition: bool,
 
     type_def_lifetime_params: DefIdMap<usize>,
 }
@@ -116,6 +117,7 @@ pub fn lower_crate(sess: &Session,
         impl_items: BTreeMap::new(),
         bodies: FxHashMap(),
         loop_scopes: Vec::new(),
+        is_in_loop_condition: false,
         type_def_lifetime_params: DefIdMap(),
     }.lower_crate(krate)
 }
@@ -251,21 +253,49 @@ impl<'a> LoweringContext<'a> {
     fn with_loop_scope<T, F>(&mut self, loop_id: NodeId, f: F) -> T
         where F: FnOnce(&mut LoweringContext) -> T
     {
+        // We're no longer in the base loop's condition; we're in another loop.
+        let was_in_loop_condition = self.is_in_loop_condition;
+        self.is_in_loop_condition = false;
+
         let len = self.loop_scopes.len();
         self.loop_scopes.push(loop_id);
+
         let result = f(self);
         assert_eq!(len + 1, self.loop_scopes.len(),
             "Loop scopes should be added and removed in stack order");
+
         self.loop_scopes.pop().unwrap();
+
+        self.is_in_loop_condition = was_in_loop_condition;
+
+        result
+    }
+
+    fn with_loop_condition_scope<T, F>(&mut self, f: F) -> T
+        where F: FnOnce(&mut LoweringContext) -> T
+    {
+        let was_in_loop_condition = self.is_in_loop_condition;
+        self.is_in_loop_condition = true;
+
+        let result = f(self);
+
+        self.is_in_loop_condition = was_in_loop_condition;
+
         result
     }
 
     fn with_new_loop_scopes<T, F>(&mut self, f: F) -> T
         where F: FnOnce(&mut LoweringContext) -> T
     {
+        let was_in_loop_condition = self.is_in_loop_condition;
+        self.is_in_loop_condition = false;
+
         let loop_scopes = mem::replace(&mut self.loop_scopes, Vec::new());
         let result = f(self);
         mem::replace(&mut self.loop_scopes, loop_scopes);
+
+        self.is_in_loop_condition = was_in_loop_condition;
+
         result
     }
 
@@ -300,17 +330,16 @@ impl<'a> LoweringContext<'a> {
         match label {
             Some((id, label_ident)) => hir::Label {
                 ident: Some(label_ident),
-                loop_id: match self.expect_full_def(id) {
-                    Def::Label(loop_id) => loop_id,
-                    _ => DUMMY_NODE_ID
+                loop_id: if let Def::Label(loop_id) = self.expect_full_def(id) {
+                    hir::LoopIdResult::Ok(loop_id)
+                } else {
+                    hir::LoopIdResult::Err(hir::LoopIdError::UnresolvedLabel)
                 }
             },
             None => hir::Label {
                 ident: None,
-                loop_id: match self.loop_scopes.last() {
-                    Some(innermost_loop_id) => *innermost_loop_id,
-                    _ => DUMMY_NODE_ID
-                }
+                loop_id: self.loop_scopes.last().map(|innermost_loop_id| Ok(*innermost_loop_id))
+                            .unwrap_or(Err(hir::LoopIdError::OutsideLoopScope)).into()
             }
         }
     }
@@ -1597,7 +1626,7 @@ impl<'a> LoweringContext<'a> {
                 ExprKind::While(ref cond, ref body, opt_ident) => {
                     self.with_loop_scope(e.id, |this|
                         hir::ExprWhile(
-                            P(this.lower_expr(cond)),
+                            this.with_loop_condition_scope(|this| P(this.lower_expr(cond))),
                             this.lower_block(body),
                             this.lower_opt_sp_ident(opt_ident)))
                 }
@@ -1699,13 +1728,29 @@ impl<'a> LoweringContext<'a> {
                     hir::ExprPath(self.lower_qpath(e.id, qself, path, ParamMode::Optional))
                 }
                 ExprKind::Break(opt_ident, ref opt_expr) => {
+                    let label_result = if self.is_in_loop_condition && opt_ident.is_none() {
+                        hir::Label {
+                            ident: opt_ident,
+                            loop_id: Err(hir::LoopIdError::UnlabeledCfInWhileCondition).into(),
+                        }
+                    } else {
+                        self.lower_label(opt_ident.map(|ident| (e.id, ident)))
+                    };
                     hir::ExprBreak(
-                        self.lower_label(opt_ident.map(|ident| (e.id, ident))),
-                                   opt_expr.as_ref().map(|x| P(self.lower_expr(x))))
+                            label_result,
+                            opt_expr.as_ref().map(|x| P(self.lower_expr(x))))
                 }
                 ExprKind::Continue(opt_ident) =>
                     hir::ExprAgain(
-                        self.lower_label(opt_ident.map(|ident| (e.id, ident)))),
+                        if self.is_in_loop_condition && opt_ident.is_none() {
+                            hir::Label {
+                                ident: opt_ident,
+                                loop_id: Err(
+                                    hir::LoopIdError::UnlabeledCfInWhileCondition).into(),
+                            }
+                        } else {
+                            self.lower_label(opt_ident.map( | ident| (e.id, ident)))
+                        }),
                 ExprKind::Ret(ref e) => hir::ExprRet(e.as_ref().map(|x| P(self.lower_expr(x)))),
                 ExprKind::InlineAsm(ref asm) => {
                     let hir_asm = hir::InlineAsm {
@@ -1846,10 +1891,12 @@ impl<'a> LoweringContext<'a> {
                     //     }
                     //   }
 
+                    // Note that the block AND the condition are evaluated in the loop scope.
+                    // This is done to allow `break` from inside the condition of the loop.
                     let (body, break_expr, sub_expr) = self.with_loop_scope(e.id, |this| (
                         this.lower_block(body),
                         this.expr_break(e.span, ThinVec::new()),
-                        P(this.lower_expr(sub_expr)),
+                        this.with_loop_condition_scope(|this| P(this.lower_expr(sub_expr))),
                     ));
 
                     // `<pat> => <body>`

src/librustc/hir/mod.rs

@@ -1030,11 +1030,56 @@ pub enum LoopSource {
     ForLoop,
 }
 
+#[derive(Clone, PartialEq, Eq, RustcEncodable, RustcDecodable, Hash, Debug, Copy)]
+pub enum LoopIdError {
+    OutsideLoopScope,
+    UnlabeledCfInWhileCondition,
+    UnresolvedLabel,
+}
+
+impl fmt::Display for LoopIdError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        fmt::Display::fmt(match *self {
+            LoopIdError::OutsideLoopScope => "not inside loop scope",
+            LoopIdError::UnlabeledCfInWhileCondition =>
+                "unlabeled control flow (break or continue) in while condition",
+            LoopIdError::UnresolvedLabel => "label not found",
+        }, f)
+    }
+}
+
+// FIXME(cramertj) this should use `Result` once master compiles w/ a vesion of Rust where
+// `Result` implements `Encodable`/`Decodable`
+#[derive(Clone, PartialEq, Eq, RustcEncodable, RustcDecodable, Hash, Debug, Copy)]
+pub enum LoopIdResult {
+    Ok(NodeId),
+    Err(LoopIdError),
+}
+impl Into<Result<NodeId, LoopIdError>> for LoopIdResult {
+    fn into(self) -> Result<NodeId, LoopIdError> {
+        match self {
+            LoopIdResult::Ok(ok) => Ok(ok),
+            LoopIdResult::Err(err) => Err(err),
+        }
+    }
+}
+impl From<Result<NodeId, LoopIdError>> for LoopIdResult {
+    fn from(res: Result<NodeId, LoopIdError>) -> Self {
+        match res {
+            Ok(ok) => LoopIdResult::Ok(ok),
+            Err(err) => LoopIdResult::Err(err),
+        }
+    }
+}
 
 #[derive(Clone, PartialEq, Eq, RustcEncodable, RustcDecodable, Hash, Debug, Copy)]
 pub struct Label {
+    // This is `Some(_)` iff there is an explicit user-specified `label
     pub ident: Option<Spanned<Ident>>,
-    pub loop_id: NodeId
+
+    // These errors are caught and then reported during the diagnostics pass in
+    // librustc_passes/loops.rs
+    pub loop_id: LoopIdResult,
 }
 
 #[derive(Clone, PartialEq, Eq, RustcEncodable, RustcDecodable, Hash, Debug, Copy)]

src/librustc/middle/liveness.rs

@@ -1003,7 +1003,10 @@ impl<'a, 'tcx> Liveness<'a, 'tcx> {
 
           hir::ExprBreak(label, ref opt_expr) => {
               // Find which label this break jumps to
-              let sc = label.loop_id;
+              let sc = match label.loop_id.into() {
+                  Ok(loop_id) => loop_id,
+                  Err(err) => span_bug!(expr.span, "loop scope error: {}", err),
+              };
 
               // Now that we know the label we're going to,
               // look it up in the break loop nodes table
@@ -1016,7 +1019,11 @@ impl<'a, 'tcx> Liveness<'a, 'tcx> {
 
           hir::ExprAgain(label) => {
               // Find which label this expr continues to
-              let sc = label.loop_id;
+              let sc = match label.loop_id.into() {
+                  Ok(loop_id) => loop_id,
+                  Err(err) => span_bug!(expr.span, "loop scope error: {}", err),
+              };
+
 
               // Now that we know the label we're going to,
               // look it up in the continue loop nodes table
@@ -1280,12 +1287,13 @@ impl<'a, 'tcx> Liveness<'a, 'tcx> {
         debug!("propagate_through_loop: using id for loop body {} {}",
                expr.id, self.ir.tcx.hir.node_to_pretty_string(body.id));
 
-        let cond_ln = match kind {
-            LoopLoop => ln,
-            WhileLoop(ref cond) => self.propagate_through_expr(&cond, ln),
-        };
-        let body_ln = self.with_loop_nodes(expr.id, succ, ln, |this| {
-            this.propagate_through_block(body, cond_ln)
+        let (cond_ln, body_ln) = self.with_loop_nodes(expr.id, succ, ln, |this| {
+            let cond_ln = match kind {
+                LoopLoop => ln,
+                WhileLoop(ref cond) => this.propagate_through_expr(&cond, ln),
+            };
+            let body_ln = this.propagate_through_block(body, cond_ln);
+            (cond_ln, body_ln)
         });
 
         // repeat until fixed point is reached: