Handle Continue correctly

Author: eholk

compiler/rustc_ast/src/ast.rs

@@ -54,7 +54,7 @@ mod tests;
 /// ```
 ///
 /// `'outer` is a label.
-#[derive(Clone, Encodable, Decodable, Copy, HashStable_Generic)]
+#[derive(Clone, Encodable, Decodable, Copy, HashStable_Generic, Eq, PartialEq)]
 pub struct Label {
     pub ident: Ident,
 }

compiler/rustc_typeck/src/check/generator_interior/drop_ranges/cfg_build.rs

@@ -4,7 +4,7 @@ use super::{
 };
 use hir::{
     intravisit::{self, Visitor},
-    Body, Expr, ExprKind, Guard, HirId,
+    Body, Expr, ExprKind, Guard, HirId, LoopIdError,
 };
 use rustc_data_structures::fx::FxHashMap;
 use rustc_hir as hir;
@@ -85,6 +85,7 @@ struct DropRangeVisitor<'a, 'tcx> {
     expr_index: PostOrderId,
     tcx: TyCtxt<'tcx>,
     typeck_results: &'a TypeckResults<'tcx>,
+    label_stack: Vec<(Option<rustc_ast::Label>, PostOrderId)>,
 }
 
 impl<'a, 'tcx> DropRangeVisitor<'a, 'tcx> {
@@ -101,7 +102,15 @@ impl<'a, 'tcx> DropRangeVisitor<'a, 'tcx> {
             hir,
             num_exprs,
         );
-        Self { hir, places, drop_ranges, expr_index: PostOrderId::from_u32(0), typeck_results, tcx }
+        Self {
+            hir,
+            places,
+            drop_ranges,
+            expr_index: PostOrderId::from_u32(0),
+            typeck_results,
+            tcx,
+            label_stack: vec![],
+        }
     }
 
     fn record_drop(&mut self, value: TrackedValue) {
@@ -210,46 +219,61 @@ impl<'a, 'tcx> DropRangeVisitor<'a, 'tcx> {
         }
     }
 
-    /// Break and continue expression targets might be another expression or a block,
-    /// but this analysis only looks at expressions. In case a break has a block as a
-    /// target, this will find the last expression in the block and return its HirId
-    /// instead.
-    fn find_target_expression(&self, hir_id: HirId) -> HirId {
-        let node = self.hir.get(hir_id);
-        match node {
-            hir::Node::Expr(_) => hir_id,
-            hir::Node::Block(b) => b.expr.map_or_else(
-                // If there is no tail expression, there will be at least one statement in the
-                // block because the block contains a break or continue statement.
-                || b.stmts.last().unwrap().hir_id,
-                |expr| expr.hir_id,
-            ),
-            hir::Node::Param(..)
-            | hir::Node::Item(..)
-            | hir::Node::ForeignItem(..)
-            | hir::Node::TraitItem(..)
-            | hir::Node::ImplItem(..)
-            | hir::Node::Variant(..)
-            | hir::Node::Field(..)
-            | hir::Node::AnonConst(..)
-            | hir::Node::Stmt(..)
-            | hir::Node::PathSegment(..)
-            | hir::Node::Ty(..)
-            | hir::Node::TraitRef(..)
-            | hir::Node::Binding(..)
-            | hir::Node::Pat(..)
-            | hir::Node::Arm(..)
-            | hir::Node::Local(..)
-            | hir::Node::Ctor(..)
-            | hir::Node::Lifetime(..)
-            | hir::Node::GenericParam(..)
-            | hir::Node::Visibility(..)
-            | hir::Node::Crate(..)
-            | hir::Node::Infer(..) => bug!("Unsupported branch target: {:?}", node),
-        }
+    /// Map a Destination to an equivalent expression node
+    ///
+    /// The destination field of a Break or Continue expression can target either an
+    /// expression or a block. The drop range analysis, however, only deals in
+    /// expression nodes, so blocks that might be the destination of a Break or Continue
+    /// will not have a PostOrderId.
+    ///
+    /// If the destination is an expression, this function will simply return that expression's
+    /// hir_id. If the destination is a block, this function will return the hir_id of last
+    /// expression in the block.
+    fn find_target_expression_from_destination(
+        &self,
+        destination: hir::Destination,
+    ) -> Result<HirId, LoopIdError> {
+        destination.target_id.map(|target| {
+            let node = self.hir.get(target);
+            match node {
+                hir::Node::Expr(_) => target,
+                hir::Node::Block(b) => find_last_block_expression(b),
+                hir::Node::Param(..)
+                | hir::Node::Item(..)
+                | hir::Node::ForeignItem(..)
+                | hir::Node::TraitItem(..)
+                | hir::Node::ImplItem(..)
+                | hir::Node::Variant(..)
+                | hir::Node::Field(..)
+                | hir::Node::AnonConst(..)
+                | hir::Node::Stmt(..)
+                | hir::Node::PathSegment(..)
+                | hir::Node::Ty(..)
+                | hir::Node::TraitRef(..)
+                | hir::Node::Binding(..)
+                | hir::Node::Pat(..)
+                | hir::Node::Arm(..)
+                | hir::Node::Local(..)
+                | hir::Node::Ctor(..)
+                | hir::Node::Lifetime(..)
+                | hir::Node::GenericParam(..)
+                | hir::Node::Visibility(..)
+                | hir::Node::Crate(..)
+                | hir::Node::Infer(..) => bug!("Unsupported branch target: {:?}", node),
+            }
+        })
     }
 }
 
+fn find_last_block_expression(block: &hir::Block<'_>) -> HirId {
+    block.expr.map_or_else(
+        // If there is no tail expression, there will be at least one statement in the
+        // block because the block contains a break or continue statement.
+        || block.stmts.last().unwrap().hir_id,
+        |expr| expr.hir_id,
+    )
+}
+
 impl<'a, 'tcx> Visitor<'tcx> for DropRangeVisitor<'a, 'tcx> {
     fn visit_expr(&mut self, expr: &'tcx Expr<'tcx>) {
         let mut reinit = None;
@@ -359,8 +383,9 @@ impl<'a, 'tcx> Visitor<'tcx> for DropRangeVisitor<'a, 'tcx> {
                 });
             }
 
-            ExprKind::Loop(body, ..) => {
+            ExprKind::Loop(body, label, ..) => {
                 let loop_begin = self.expr_index + 1;
+                self.label_stack.push((label, loop_begin));
                 if body.stmts.is_empty() && body.expr.is_none() {
                     // For empty loops we won't have updated self.expr_index after visiting the
                     // body, meaning we'd get an edge from expr_index to expr_index + 1, but
@@ -370,11 +395,31 @@ impl<'a, 'tcx> Visitor<'tcx> for DropRangeVisitor<'a, 'tcx> {
                     self.visit_block(body);
                     self.drop_ranges.add_control_edge(self.expr_index, loop_begin);
                 }
+                self.label_stack.pop();
             }
-            ExprKind::Break(hir::Destination { target_id: Ok(target), .. }, ..)
-            | ExprKind::Continue(hir::Destination { target_id: Ok(target), .. }, ..) => {
-                self.drop_ranges
-                    .add_control_edge_hir_id(self.expr_index, self.find_target_expression(target));
+            // Find the loop entry by searching through the label stack for either the last entry
+            // (if label is none), or the first entry where the label matches this one. The Loop
+            // case maintains this stack mapping labels to the PostOrderId for the loop entry.
+            ExprKind::Continue(hir::Destination { label, .. }, ..) => self
+                .label_stack
+                .iter()
+                .rev()
+                .find(|(loop_label, _)| label.is_none() || *loop_label == label)
+                .map_or((), |(_, target)| {
+                    self.drop_ranges.add_control_edge(self.expr_index, *target)
+                }),
+
+            ExprKind::Break(destination, ..) => {
+                // destination either points to an expression or to a block. We use
+                // find_target_expression_from_destination to use the last expression of the block
+                // if destination points to a block.
+                //
+                // We add an edge to the hir_id of the expression/block we are breaking out of, and
+                // then in process_deferred_edges we will map this hir_id to its PostOrderId, which
+                // will refer to the end of the block due to the post order traversal.
+                self.find_target_expression_from_destination(destination).map_or((), |target| {
+                    self.drop_ranges.add_control_edge_hir_id(self.expr_index, target)
+                })
             }
 
             ExprKind::Call(f, args) => {
@@ -399,11 +444,9 @@ impl<'a, 'tcx> Visitor<'tcx> for DropRangeVisitor<'a, 'tcx> {
             | ExprKind::Binary(..)
             | ExprKind::Block(..)
             | ExprKind::Box(..)
-            | ExprKind::Break(..)
             | ExprKind::Cast(..)
             | ExprKind::Closure(..)
             | ExprKind::ConstBlock(..)
-            | ExprKind::Continue(..)
             | ExprKind::DropTemps(..)
             | ExprKind::Err
             | ExprKind::Field(..)

src/test/ui/generator/drop-control-flow.rs

@@ -110,6 +110,22 @@ fn nested_loop() {
     };
 }
 
+fn loop_continue(b: bool) {
+    let _ = || {
+        let mut arr = [Ptr];
+        let mut count = 0;
+        drop(arr);
+        while count < 3 {
+            count += 1;
+            yield;
+            if b {
+                arr = [Ptr];
+                continue;
+            }
+        }
+    };
+}
+
 fn main() {
     one_armed_if(true);
     if_let(Some(41));
@@ -118,4 +134,5 @@ fn main() {
     reinit();
     loop_uninit();
     nested_loop();
+    loop_continue(true);
 }