Generator drop tracking: improve break and continue handling

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_interface/src/tests.rs

@@ -730,6 +730,7 @@ fn test_debugging_options_tracking_hash() {
     tracked!(debug_info_for_profiling, true);
     tracked!(debug_macros, true);
     tracked!(dep_info_omit_d_target, true);
+    tracked!(drop_tracking, true);
     tracked!(dual_proc_macros, true);
     tracked!(fewer_names, Some(true));
     tracked!(force_unstable_if_unmarked, true);

compiler/rustc_session/src/options.rs

@@ -1173,6 +1173,8 @@ options! {
     dont_buffer_diagnostics: bool = (false, parse_bool, [UNTRACKED],
         "emit diagnostics rather than buffering (breaks NLL error downgrading, sorting) \
         (default: no)"),
+    drop_tracking: bool = (false, parse_bool, [TRACKED],
+        "enables drop tracking in generators (default: no)"),
     dual_proc_macros: bool = (false, parse_bool, [TRACKED],
         "load proc macros for both target and host, but only link to the target (default: no)"),
     dump_dep_graph: bool = (false, parse_bool, [UNTRACKED],

compiler/rustc_typeck/src/check/generator_interior.rs

@@ -22,11 +22,6 @@ use tracing::debug;
 
 mod drop_ranges;
 
-// FIXME(eholk): This flag is here to give a quick way to disable drop tracking in case we find
-// unexpected breakages while it's still new. It should be removed before too long. For example,
-// see #93161.
-const ENABLE_DROP_TRACKING: bool = false;
-
 struct InteriorVisitor<'a, 'tcx> {
     fcx: &'a FnCtxt<'a, 'tcx>,
     types: FxIndexSet<ty::GeneratorInteriorTypeCause<'tcx>>,
@@ -82,7 +77,7 @@ impl<'a, 'tcx> InteriorVisitor<'a, 'tcx> {
                                 yield_data.expr_and_pat_count, self.expr_count, source_span
                             );
 
-                            if ENABLE_DROP_TRACKING
+                            if self.fcx.sess().opts.debugging_opts.drop_tracking
                                 && self
                                     .drop_ranges
                                     .is_dropped_at(hir_id, yield_data.expr_and_pat_count)

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

@@ -37,7 +37,7 @@ pub fn compute_drop_ranges<'a, 'tcx>(
     def_id: DefId,
     body: &'tcx Body<'tcx>,
 ) -> DropRanges {
-    if super::ENABLE_DROP_TRACKING {
+    if fcx.sess().opts.debugging_opts.drop_tracking {
         let consumed_borrowed_places = find_consumed_and_borrowed(fcx, def_id, body);
 
         let num_exprs = fcx.tcx.region_scope_tree(def_id).body_expr_count(body.id()).unwrap_or(0);

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) {
@@ -209,6 +218,60 @@ impl<'a, 'tcx> DropRangeVisitor<'a, 'tcx> {
             self.drop_ranges.add_control_edge(self.expr_index + 1, self.expr_index + 1);
         }
     }
+
+    /// 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> {
@@ -320,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
@@ -331,10 +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, 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) => {
@@ -359,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(..)
@@ -462,11 +545,13 @@ impl DropRangesBuilder {
     /// Should be called after visiting the HIR but before solving the control flow, otherwise some
     /// edges will be missed.
     fn process_deferred_edges(&mut self) {
+        trace!("processing deferred edges. post_order_map={:#?}", self.post_order_map);
         let mut edges = vec![];
         swap(&mut edges, &mut self.deferred_edges);
         edges.into_iter().for_each(|(from, to)| {
-            let to = *self.post_order_map.get(&to).expect("Expression ID not found");
             trace!("Adding deferred edge from {:?} to {:?}", from, to);
+            let to = *self.post_order_map.get(&to).expect("Expression ID not found");
+            trace!("target edge PostOrderId={:?}", to);
             self.add_control_edge(from, to)
         });
     }

src/test/ui/async-await/issue-93197.rs

@@ -1,6 +1,7 @@
 // Regression test for #93197
 // check-pass
 // edition:2021
+// compile-flags: -Zdrop-tracking
 
 #![feature(try_blocks)]
 

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

@@ -1,4 +1,5 @@
 // build-pass
+// compile-flags: -Zdrop-tracking
 
 // FIXME(eholk): temporarily disabled while drop range tracking is disabled
 // (see generator_interior.rs:27)
@@ -114,6 +115,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));
@@ -122,4 +139,5 @@ fn main() {
     reinit();
     loop_uninit();
     nested_loop();
+    loop_continue(true);
 }