Fix control flow handling in generator_interior

All tests pass now! The issue was that we weren't handling all edges
correctly, but now they are handled consistently.

This includes code to dump a graphviz file for the CFG we built for drop
tracking.

Also removes old DropRanges tests.

5 files changed, 78 insertions, 22 deletions

diff --git a/Cargo.lock b/Cargo.lock
index b46696dde7c..cde44f96ce7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4388,6 +4388,7 @@ dependencies = [
  "rustc_attr",
  "rustc_data_structures",
  "rustc_errors",
+ "rustc_graphviz",
  "rustc_hir",
  "rustc_hir_pretty",
  "rustc_index",
diff --git a/compiler/rustc_typeck/Cargo.toml b/compiler/rustc_typeck/Cargo.toml
index 3279c331ade..57930a28a35 100644
--- a/compiler/rustc_typeck/Cargo.toml
+++ b/compiler/rustc_typeck/Cargo.toml
@@ -15,6 +15,7 @@ rustc_middle = { path = "../rustc_middle" }
 rustc_attr = { path = "../rustc_attr" }
 rustc_data_structures = { path = "../rustc_data_structures" }
 rustc_errors = { path = "../rustc_errors" }
+rustc_graphviz = { path = "../rustc_graphviz" }
 rustc_hir = { path = "../rustc_hir" }
 rustc_hir_pretty = { path = "../rustc_hir_pretty" }
 rustc_target = { path = "../rustc_target" }
diff --git a/compiler/rustc_typeck/src/check/generator_interior.rs b/compiler/rustc_typeck/src/check/generator_interior.rs
index 1df45b566ce..f26ba875c01 100644
--- a/compiler/rustc_typeck/src/check/generator_interior.rs
+++ b/compiler/rustc_typeck/src/check/generator_interior.rs
@@ -26,9 +26,6 @@ use rustc_span::Span;
 use smallvec::SmallVec;
 use tracing::debug;
 
-#[cfg(test)]
-mod tests;
-
 mod drop_ranges;
 
 struct InteriorVisitor<'a, 'tcx> {
@@ -255,6 +252,7 @@ pub fn resolve_interior<'a, 'tcx>(
         intravisit::walk_body(&mut drop_range_visitor, body);
 
         drop_range_visitor.drop_ranges.propagate_to_fixpoint();
+        // drop_range_visitor.drop_ranges.save_graph("drop_ranges.dot");
 
         InteriorVisitor {
             fcx,
@@ -877,18 +875,18 @@ impl<'tcx> Visitor<'tcx> for DropRangeVisitor<'tcx> {
                 reinit = Some(lhs);
             }
             ExprKind::Loop(body, ..) => {
-                let loop_begin = self.expr_count;
+                let loop_begin = self.expr_count + 1;
                 self.visit_block(body);
                 self.drop_ranges.add_control_edge(self.expr_count, loop_begin);
             }
             ExprKind::Match(scrutinee, arms, ..) => {
                 self.visit_expr(scrutinee);
 
-                let fork = self.expr_count - 1;
+                let fork = self.expr_count;
                 let arm_end_ids = arms
                     .iter()
                     .map(|Arm { pat, body, guard, .. }| {
-                        self.drop_ranges.add_control_edge(fork, self.expr_count);
+                        self.drop_ranges.add_control_edge(fork, self.expr_count + 1);
                         self.visit_pat(pat);
                         match guard {
                             Some(Guard::If(expr)) => self.visit_expr(expr),
@@ -914,8 +912,8 @@ impl<'tcx> Visitor<'tcx> for DropRangeVisitor<'tcx> {
             _ => intravisit::walk_expr(self, expr),
         }
 
-        self.drop_ranges.add_node_mapping(expr.hir_id, self.expr_count);
         self.expr_count += 1;
+        self.drop_ranges.add_node_mapping(expr.hir_id, self.expr_count);
         self.consume_expr(expr);
         if let Some(expr) = reinit {
             self.reinit_expr(expr);
diff --git a/compiler/rustc_typeck/src/check/generator_interior/drop_ranges.rs b/compiler/rustc_typeck/src/check/generator_interior/drop_ranges.rs
index 504734080b6..5fe3e408838 100644
--- a/compiler/rustc_typeck/src/check/generator_interior/drop_ranges.rs
+++ b/compiler/rustc_typeck/src/check/generator_interior/drop_ranges.rs
@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
 use std::fmt::Debug;
 use std::mem::swap;
 
+use rustc_graphviz as dot;
 use rustc_hir::{HirId, HirIdMap};
 use rustc_index::bit_set::BitSet;
 use rustc_index::vec::IndexVec;
@@ -182,7 +183,9 @@ impl DropRanges {
             changed
         };
 
-        while propagate() {}
+        while propagate() {
+            trace!("drop_state changed, re-running propagation");
+        }
 
         trace!("after fixpoint: {:#?}", self);
     }
@@ -200,6 +203,73 @@ impl DropRanges {
         }
         preds
     }
+
+    // pub fn save_graph(&self, filename: &str) {
+    //     use std::fs::File;
+    //     dot::render(self, &mut File::create(filename).unwrap()).unwrap();
+    // }
+}
+
+impl<'a> dot::GraphWalk<'a> for DropRanges {
+    type Node = PostOrderId;
+
+    type Edge = (PostOrderId, PostOrderId);
+
+    fn nodes(&'a self) -> dot::Nodes<'a, Self::Node> {
+        self.nodes.iter_enumerated().map(|(i, _)| i).collect()
+    }
+
+    fn edges(&'a self) -> dot::Edges<'a, Self::Edge> {
+        self.nodes
+            .iter_enumerated()
+            .flat_map(|(i, node)| {
+                if node.successors.len() == 0 {
+                    vec![(i, PostOrderId::from_usize(i.index() + 1))]
+                } else {
+                    node.successors.iter().map(move |&s| (i, s)).collect()
+                }
+            })
+            .collect()
+    }
+
+    fn source(&'a self, edge: &Self::Edge) -> Self::Node {
+        edge.0
+    }
+
+    fn target(&'a self, edge: &Self::Edge) -> Self::Node {
+        edge.1
+    }
+}
+
+impl<'a> dot::Labeller<'a> for DropRanges {
+    type Node = PostOrderId;
+
+    type Edge = (PostOrderId, PostOrderId);
+
+    fn graph_id(&'a self) -> dot::Id<'a> {
+        dot::Id::new("drop_ranges").unwrap()
+    }
+
+    fn node_id(&'a self, n: &Self::Node) -> dot::Id<'a> {
+        dot::Id::new(format!("id{}", n.index())).unwrap()
+    }
+
+    fn node_label(&'a self, n: &Self::Node) -> dot::LabelText<'a> {
+        dot::LabelText::LabelStr(
+            format!(
+                "{:?}, local_id: {}",
+                n,
+                self.post_order_map
+                    .iter()
+                    .find(|(_hir_id, &post_order_id)| post_order_id == n.index())
+                    .map_or("<unknown>".into(), |(hir_id, _)| format!(
+                        "{}",
+                        hir_id.local_id.index()
+                    ))
+            )
+            .into(),
+        )
+    }
 }
 
 #[derive(Debug)]
diff --git a/compiler/rustc_typeck/src/check/generator_interior/tests.rs b/compiler/rustc_typeck/src/check/generator_interior/tests.rs
deleted file mode 100644
index 8f973bb9489..00000000000
--- a/compiler/rustc_typeck/src/check/generator_interior/tests.rs
+++ /dev/null
@@ -1,14 +0,0 @@
-use super::DropRange;
-
-#[test]
-fn drop_range_uses_last_event() {
-    let mut range = DropRange::empty();
-    range.drop(10);
-    range.reinit(10);
-    assert!(!range.is_dropped_at(10));
-
-    let mut range = DropRange::empty();
-    range.reinit(10);
-    range.drop(10);
-    assert!(range.is_dropped_at(10));
-}