Rollup merge of #106699 - eholk:await-chains-drop-tracking, r=wesleywiser

[drop tracking] Visit break expressions

This fixes #102383 by remembering to visit the expression in `break expr` when building the drop tracking CFG. Missing this step was causing an off-by-one error which meant after a number of awaits we'd be
looking for dropped values at the wrong point in the code.

Additionally, this changes the order of traversal for assignment expressions to visit the rhs and then the lhs. This matches what is done elsewhere.

Finally, this improves some of the debugging output (for example, the CFG visualizer) to make it easier to figure out these sorts of issues.

Author: matthiaskrgr

compiler/rustc_hir_typeck/src/generator_interior/drop_ranges/cfg_build.rs

@@ -304,8 +304,8 @@ impl<'a, 'tcx> Visitor<'tcx> for DropRangeVisitor<'a, 'tcx> {
         let mut reinit = None;
         match expr.kind {
             ExprKind::Assign(lhs, rhs, _) => {
-                self.visit_expr(lhs);
                 self.visit_expr(rhs);
+                self.visit_expr(lhs);
 
                 reinit = Some(lhs);
             }
@@ -433,7 +433,7 @@ impl<'a, 'tcx> Visitor<'tcx> for DropRangeVisitor<'a, 'tcx> {
                     self.drop_ranges.add_control_edge(self.expr_index, *target)
                 }),
 
-            ExprKind::Break(destination, ..) => {
+            ExprKind::Break(destination, value) => {
                 // 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.
@@ -443,7 +443,11 @@ impl<'a, 'tcx> Visitor<'tcx> for DropRangeVisitor<'a, 'tcx> {
                 // 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)
-                })
+                });
+
+                if let Some(value) = value {
+                    self.visit_expr(value);
+                }
             }
 
             ExprKind::Call(f, args) => {
@@ -465,6 +469,12 @@ impl<'a, 'tcx> Visitor<'tcx> for DropRangeVisitor<'a, 'tcx> {
 
             ExprKind::AddrOf(..)
             | ExprKind::Array(..)
+            // FIXME(eholk): We probably need special handling for AssignOps. The ScopeTree builder
+            // in region.rs runs both lhs then rhs and rhs then lhs and then sets all yields to be
+            // the latest they show up in either traversal. With the older scope-based
+            // approximation, this was fine, but it's probably not right now. What we probably want
+            // to do instead is still run both orders, but consider anything that showed up as a
+            // yield in either order.
             | ExprKind::AssignOp(..)
             | ExprKind::Binary(..)
             | ExprKind::Block(..)
@@ -502,6 +512,9 @@ impl<'a, 'tcx> Visitor<'tcx> for DropRangeVisitor<'a, 'tcx> {
 
         // Increment expr_count here to match what InteriorVisitor expects.
         self.expr_index = self.expr_index + 1;
+
+        // Save a node mapping to get better CFG visualization
+        self.drop_ranges.add_node_mapping(pat.hir_id, self.expr_index);
     }
 }
 
@@ -521,7 +534,7 @@ impl DropRangesBuilder {
                 }
             });
         }
-        debug!("hir_id_map: {:?}", tracked_value_map);
+        debug!("hir_id_map: {:#?}", tracked_value_map);
         let num_values = tracked_value_map.len();
         Self {
             tracked_value_map,

compiler/rustc_hir_typeck/src/generator_interior/drop_ranges/cfg_visualize.rs

@@ -2,6 +2,7 @@
 //! flow graph when needed for debugging.
 
 use rustc_graphviz as dot;
+use rustc_hir::{Expr, ExprKind, Node};
 use rustc_middle::ty::TyCtxt;
 
 use super::{DropRangesBuilder, PostOrderId};
@@ -80,10 +81,14 @@ impl<'a> dot::Labeller<'a> for DropRangesGraph<'_, '_> {
                     .post_order_map
                     .iter()
                     .find(|(_hir_id, &post_order_id)| post_order_id == *n)
-                    .map_or("<unknown>".into(), |(hir_id, _)| self
-                        .tcx
-                        .hir()
-                        .node_to_string(*hir_id))
+                    .map_or("<unknown>".into(), |(hir_id, _)| format!(
+                        "{}{}",
+                        self.tcx.hir().node_to_string(*hir_id),
+                        match self.tcx.hir().find(*hir_id) {
+                            Some(Node::Expr(Expr { kind: ExprKind::Yield(..), .. })) => " (yield)",
+                            _ => "",
+                        }
+                    ))
             )
             .into(),
         )

compiler/rustc_hir_typeck/src/generator_interior/mod.rs

@@ -71,10 +71,8 @@ impl<'a, 'tcx> InteriorVisitor<'a, 'tcx> {
                                 yield_data.expr_and_pat_count, self.expr_count, source_span
                             );
 
-                            if self.fcx.sess().opts.unstable_opts.drop_tracking
-                                && self
-                                    .drop_ranges
-                                    .is_dropped_at(hir_id, yield_data.expr_and_pat_count)
+                            if self
+                                .is_dropped_at_yield_location(hir_id, yield_data.expr_and_pat_count)
                             {
                                 debug!("value is dropped at yield point; not recording");
                                 return false;
@@ -173,6 +171,18 @@ impl<'a, 'tcx> InteriorVisitor<'a, 'tcx> {
             }
         }
     }
+
+    /// If drop tracking is enabled, consult drop_ranges to see if a value is
+    /// known to be dropped at a yield point and therefore can be omitted from
+    /// the generator witness.
+    fn is_dropped_at_yield_location(&self, value_hir_id: HirId, yield_location: usize) -> bool {
+        // short-circuit if drop tracking is not enabled.
+        if !self.fcx.sess().opts.unstable_opts.drop_tracking {
+            return false;
+        }
+
+        self.drop_ranges.is_dropped_at(value_hir_id, yield_location)
+    }
 }
 
 pub fn resolve_interior<'a, 'tcx>(

tests/ui/async-await/await-sequence.rs

@@ -0,0 +1,21 @@
+// edition:2021
+// compile-flags: -Z drop-tracking
+// build-pass
+
+use std::collections::HashMap;
+
+fn main() {
+    let _ = real_main();
+}
+
+async fn nop() {}
+
+async fn real_main() {
+    nop().await;
+    nop().await;
+    nop().await;
+    nop().await;
+
+    let mut map: HashMap<(), ()> = HashMap::new();
+    map.insert((), nop().await);
+}