Fix drop tracking ICEs and re-enable generator drop tracking

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_mir_transform/src/generator.rs

@@ -749,9 +749,20 @@ fn sanitize_witness<'tcx>(
         }
         let decl_ty = tcx.normalize_erasing_regions(param_env, decl.ty);
 
+        let is_uninhabited = tcx.is_ty_uninhabited_from(
+            tcx.parent_module(tcx.hir().local_def_id_to_hir_id(did.expect_local())).to_def_id(),
+            decl_ty,
+            param_env,
+        );
+
         // Sanity check that typeck knows about the type of locals which are
         // live across a suspension point
-        if !allowed.contains(&decl_ty) && !allowed_upvars.contains(&decl_ty) {
+        if is_uninhabited || allowed.contains(&decl_ty) || allowed_upvars.contains(&decl_ty) {
+            // This type which appears in the generator either...
+            // - is uninhabited, in which case it can't actually be captured at runtime
+            // - appears in the approximation from the static type (`allowed`)
+            // - appears in the list of upvars ...
+        } else {
             span_bug!(
                 body.span,
                 "Broken MIR: generator contains type {} in MIR, \

compiler/rustc_typeck/src/check/generator_interior.rs

@@ -25,7 +25,7 @@ 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;
+const ENABLE_DROP_TRACKING: bool = true;
 
 struct InteriorVisitor<'a, 'tcx> {
     fcx: &'a FnCtxt<'a, 'tcx>,

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/async-fn-nonsend.rs

@@ -1,10 +1,6 @@
 // edition:2018
 // compile-flags: --crate-type lib
 
-// FIXME(eholk): temporarily disabled while drop range tracking is disabled
-// (see generator_interior.rs:27)
-// ignore-test
-
 use std::{cell::RefCell, fmt::Debug, rc::Rc};
 
 fn non_sync() -> impl Debug {

src/test/ui/async-await/unresolved_type_param.rs

@@ -3,10 +3,6 @@
 // (rather than give a general error message)
 // edition:2018
 
-// FIXME(eholk): temporarily disabled while drop range tracking is disabled
-// (see generator_interior.rs:27)
-// ignore-test
-
 async fn bar<T>() -> () {}
 
 async fn foo() {

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

@@ -1,9 +1,5 @@
 // build-pass
 
-// FIXME(eholk): temporarily disabled while drop range tracking is disabled
-// (see generator_interior.rs:27)
-// ignore-test
-
 // A test to ensure generators capture values that were conditionally dropped,
 // and also that values that are dropped along all paths to a yield do not get
 // included in the generator type.
@@ -114,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));
@@ -122,4 +134,5 @@ fn main() {
     reinit();
     loop_uninit();
     nested_loop();
+    loop_continue(true);
 }

src/test/ui/generator/issue-57478.rs

@@ -1,9 +1,5 @@
 // check-pass
 
-// FIXME(eholk): temporarily disabled while drop range tracking is disabled
-// (see generator_interior.rs:27)
-// ignore-test
-
 #![feature(negative_impls, generators)]
 
 struct Foo;

src/test/ui/generator/issue-93161.rs

@@ -1,31 +1,8 @@
 // edition:2021
-// run-pass
+// build-pass
 
 #![feature(never_type)]
 
-use std::future::Future;
-
-// See if we can run a basic `async fn`
-pub async fn foo(x: &u32, y: u32) -> u32 {
-    let y = &y;
-    let z = 9;
-    let z = &z;
-    let y = async { *y + *z }.await;
-    let a = 10;
-    let a = &a;
-    *x + y + *a
-}
-
-async fn add(x: u32, y: u32) -> u32 {
-    let a = async { x + y };
-    a.await
-}
-
-async fn build_aggregate(a: u32, b: u32, c: u32, d: u32) -> u32 {
-    let x = (add(a, b).await, add(c, d).await);
-    x.0 + x.1
-}
-
 enum Never {}
 fn never() -> Never {
     panic!()
@@ -43,51 +20,6 @@ async fn includes_never(crash: bool, x: u32) -> u32 {
     result
 }
 
-async fn partial_init(x: u32) -> u32 {
-    #[allow(unreachable_code)]
-    let _x: (String, !) = (String::new(), return async { x + x }.await);
-}
-
-async fn read_exact(_from: &mut &[u8], _to: &mut [u8]) -> Option<()> {
-    Some(())
-}
-
-async fn hello_world() {
-    let data = [0u8; 1];
-    let mut reader = &data[..];
-
-    let mut marker = [0u8; 1];
-    read_exact(&mut reader, &mut marker).await.unwrap();
-}
-
-fn run_fut<T>(fut: impl Future<Output = T>) -> T {
-    use std::sync::Arc;
-    use std::task::{Context, Poll, Wake, Waker};
-
-    struct MyWaker;
-    impl Wake for MyWaker {
-        fn wake(self: Arc<Self>) {
-            unimplemented!()
-        }
-    }
-
-    let waker = Waker::from(Arc::new(MyWaker));
-    let mut context = Context::from_waker(&waker);
-
-    let mut pinned = Box::pin(fut);
-    loop {
-        match pinned.as_mut().poll(&mut context) {
-            Poll::Pending => continue,
-            Poll::Ready(v) => return v,
-        }
-    }
-}
-
 fn main() {
-    let x = 5;
-    assert_eq!(run_fut(foo(&x, 7)), 31);
-    assert_eq!(run_fut(build_aggregate(1, 2, 3, 4)), 10);
-    assert_eq!(run_fut(includes_never(false, 4)), 16);
-    assert_eq!(run_fut(partial_init(4)), 8);
-    run_fut(hello_world());
+    let _ = includes_never(false, 4);
 }

src/test/ui/generator/partial-drop.rs

@@ -1,7 +1,3 @@
-// FIXME(eholk): temporarily disabled while drop range tracking is disabled
-// (see generator_interior.rs:27)
-// ignore-test
-
 #![feature(negative_impls, generators)]
 
 struct Foo;