More comments and refactoring

The refactoring mainly keeps the separation between the modules clearer.
For example, process_deferred_edges function moved to cfg_build.rs since
that is really part of building the CFG, not finding the fixpoint.

Also, we use PostOrderId instead of usize in a lot more places now.

4 files changed, 117 insertions, 75 deletions

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 b200320b8d3..9fbefcfb088 100644
--- a/compiler/rustc_typeck/src/check/generator_interior/drop_ranges.rs
+++ b/compiler/rustc_typeck/src/check/generator_interior/drop_ranges.rs
@@ -41,7 +41,7 @@ pub fn compute_drop_ranges<'a, 'tcx>(
 
     drop_ranges.propagate_to_fixpoint();
 
-    drop_ranges
+    DropRanges { hir_id_map: drop_ranges.hir_id_map, nodes: drop_ranges.nodes }
 }
 
 /// Applies `f` to consumable portion of a HIR node.
@@ -77,12 +77,59 @@ rustc_index::newtype_index! {
 pub struct DropRanges {
     hir_id_map: HirIdMap<HirIdIndex>,
     nodes: IndexVec<PostOrderId, NodeInfo>,
-    deferred_edges: Vec<(usize, HirId)>,
-    // FIXME: This should only be used for loops and break/continue.
-    post_order_map: HirIdMap<usize>,
 }
 
-impl Debug for DropRanges {
+impl DropRanges {
+    pub fn is_dropped_at(&self, hir_id: HirId, location: usize) -> bool {
+        self.hir_id_map
+            .get(&hir_id)
+            .copied()
+            .map_or(false, |hir_id| self.expect_node(location.into()).drop_state.contains(hir_id))
+    }
+
+    /// Returns a reference to the NodeInfo for a node, panicking if it does not exist
+    fn expect_node(&self, id: PostOrderId) -> &NodeInfo {
+        &self.nodes[id]
+    }
+}
+
+/// Tracks information needed to compute drop ranges.
+struct DropRangesBuilder {
+    /// The core of DropRangesBuilder is a set of nodes, which each represent
+    /// one expression. We primarily refer to them by their index in a
+    /// post-order traversal of the HIR tree,  since this is what
+    /// generator_interior uses to talk about yield positions.
+    ///
+    /// This IndexVec keeps the relevant details for each node. See the
+    /// NodeInfo struct for more details, but this information includes things
+    /// such as the set of control-flow successors, which variables are dropped
+    /// or reinitialized, and whether each variable has been inferred to be
+    /// known-dropped or potentially reintiialized at each point.
+    nodes: IndexVec<PostOrderId, NodeInfo>,
+    /// We refer to values whose drop state we are tracking by the HirId of
+    /// where they are defined. Within a NodeInfo, however, we store the
+    /// drop-state in a bit vector indexed by a HirIdIndex
+    /// (see NodeInfo::drop_state). The hir_id_map field stores the mapping
+    /// from HirIds to the HirIdIndex that is used to represent that value in
+    /// bitvector.
+    hir_id_map: HirIdMap<HirIdIndex>,
+
+    /// When building the control flow graph, we don't always know the
+    /// post-order index of the target node at the point we encounter it.
+    /// For example, this happens with break and continue. In those cases,
+    /// we store a pair of the PostOrderId of the source and the HirId
+    /// of the target. Once we have gathered all of these edges, we make a
+    /// pass over the set of deferred edges (see process_deferred_edges in
+    /// cfg_build.rs), look up the PostOrderId for the target (since now the
+    /// post-order index for all nodes is known), and add missing control flow
+    /// edges.
+    deferred_edges: Vec<(PostOrderId, HirId)>,
+    /// This maps HirIds of expressions to their post-order index. It is
+    /// used in process_deferred_edges to correctly add back-edges.
+    post_order_map: HirIdMap<PostOrderId>,
+}
+
+impl Debug for DropRangesBuilder {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("DropRanges")
             .field("hir_id_map", &self.hir_id_map)
@@ -98,32 +145,20 @@ impl Debug for DropRanges {
 /// by their index in the post-order traversal. At its core, DropRanges maps
 /// (hir_id, post_order_id) -> bool, where a true value indicates that the value is definitely
 /// dropped at the point of the node identified by post_order_id.
-impl DropRanges {
-    pub fn is_dropped_at(&mut self, hir_id: HirId, location: usize) -> bool {
-        self.hir_id_map
-            .get(&hir_id)
-            .copied()
-            .map_or(false, |hir_id| self.expect_node(location.into()).drop_state.contains(hir_id))
-    }
-
+impl DropRangesBuilder {
     /// Returns the number of values (hir_ids) that are tracked
     fn num_values(&self) -> usize {
         self.hir_id_map.len()
     }
 
-    /// Returns a reference to the NodeInfo for a node, panicking if it does not exist
-    fn expect_node(&self, id: PostOrderId) -> &NodeInfo {
-        &self.nodes[id]
-    }
-
     fn node_mut(&mut self, id: PostOrderId) -> &mut NodeInfo {
         let size = self.num_values();
         self.nodes.ensure_contains_elem(id, || NodeInfo::new(size));
         &mut self.nodes[id]
     }
 
-    fn add_control_edge(&mut self, from: usize, to: usize) {
-        trace!("adding control edge from {} to {}", from, to);
+    fn add_control_edge(&mut self, from: PostOrderId, to: PostOrderId) {
+        trace!("adding control edge from {:?} to {:?}", from, to);
         self.node_mut(from.into()).successors.push(to.into());
     }
 }
diff --git a/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/cfg_build.rs b/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/cfg_build.rs
index e1f1b44283b..32f423f3bfe 100644
--- a/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/cfg_build.rs
+++ b/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/cfg_build.rs
@@ -1,6 +1,6 @@
 use super::{
-    for_each_consumable, record_consumed_borrow::ConsumedAndBorrowedPlaces, DropRanges, HirIdIndex,
-    NodeInfo,
+    for_each_consumable, record_consumed_borrow::ConsumedAndBorrowedPlaces, DropRangesBuilder,
+    HirIdIndex, NodeInfo, PostOrderId,
 };
 use hir::{
     intravisit::{self, NestedVisitorMap, Visitor},
@@ -9,20 +9,24 @@ use hir::{
 use rustc_hir as hir;
 use rustc_index::vec::IndexVec;
 use rustc_middle::hir::map::Map;
+use std::mem::swap;
 
 /// Traverses the body to find the control flow graph and locations for the
 /// relevant places are dropped or reinitialized.
 ///
 /// The resulting structure still needs to be iterated to a fixed point, which
 /// can be done with propagate_to_fixpoint in cfg_propagate.
-pub fn build_control_flow_graph<'tcx>(
+pub(super) fn build_control_flow_graph<'tcx>(
     hir: Map<'tcx>,
     consumed_borrowed_places: ConsumedAndBorrowedPlaces,
     body: &'tcx Body<'tcx>,
     num_exprs: usize,
-) -> DropRanges {
+) -> DropRangesBuilder {
     let mut drop_range_visitor = DropRangeVisitor::new(hir, consumed_borrowed_places, num_exprs);
     intravisit::walk_body(&mut drop_range_visitor, body);
+
+    drop_range_visitor.drop_ranges.process_deferred_edges();
+
     drop_range_visitor.drop_ranges
 }
 
@@ -35,27 +39,27 @@ pub fn build_control_flow_graph<'tcx>(
 struct DropRangeVisitor<'tcx> {
     hir: Map<'tcx>,
     places: ConsumedAndBorrowedPlaces,
-    drop_ranges: DropRanges,
-    expr_count: usize,
+    drop_ranges: DropRangesBuilder,
+    expr_index: PostOrderId,
 }
 
 impl<'tcx> DropRangeVisitor<'tcx> {
     fn new(hir: Map<'tcx>, places: ConsumedAndBorrowedPlaces, num_exprs: usize) -> Self {
         debug!("consumed_places: {:?}", places.consumed);
-        let drop_ranges = DropRanges::new(
+        let drop_ranges = DropRangesBuilder::new(
             places.consumed.iter().flat_map(|(_, places)| places.iter().copied()),
             hir,
             num_exprs,
         );
-        Self { hir, places, drop_ranges, expr_count: 0 }
+        Self { hir, places, drop_ranges, expr_index: PostOrderId::from_u32(0) }
     }
 
     fn record_drop(&mut self, hir_id: HirId) {
         if self.places.borrowed.contains(&hir_id) {
             debug!("not marking {:?} as dropped because it is borrowed at some point", hir_id);
         } else {
-            debug!("marking {:?} as dropped at {}", hir_id, self.expr_count);
-            let count = self.expr_count;
+            debug!("marking {:?} as dropped at {:?}", hir_id, self.expr_index);
+            let count = self.expr_index;
             self.drop_ranges.drop_at(hir_id, count);
         }
     }
@@ -63,7 +67,7 @@ impl<'tcx> DropRangeVisitor<'tcx> {
     /// ExprUseVisitor's consume callback doesn't go deep enough for our purposes in all
     /// expressions. This method consumes a little deeper into the expression when needed.
     fn consume_expr(&mut self, expr: &hir::Expr<'_>) {
-        debug!("consuming expr {:?}, count={}", expr.hir_id, self.expr_count);
+        debug!("consuming expr {:?}, count={:?}", expr.hir_id, self.expr_index);
         let places = self
             .places
             .consumed
@@ -80,8 +84,8 @@ impl<'tcx> DropRangeVisitor<'tcx> {
             hir::Path { res: hir::def::Res::Local(hir_id), .. },
         )) = expr.kind
         {
-            let location = self.expr_count;
-            debug!("reinitializing {:?} at {}", hir_id, location);
+            let location = self.expr_index;
+            debug!("reinitializing {:?} at {:?}", hir_id, location);
             self.drop_ranges.reinit_at(*hir_id, location);
         } else {
             debug!("reinitializing {:?} is not supported", expr);
@@ -102,18 +106,18 @@ impl<'tcx> Visitor<'tcx> for DropRangeVisitor<'tcx> {
             ExprKind::If(test, if_true, if_false) => {
                 self.visit_expr(test);
 
-                let fork = self.expr_count;
+                let fork = self.expr_index;
 
-                self.drop_ranges.add_control_edge(fork, self.expr_count + 1);
+                self.drop_ranges.add_control_edge(fork, self.expr_index + 1);
                 self.visit_expr(if_true);
-                let true_end = self.expr_count;
+                let true_end = self.expr_index;
 
-                self.drop_ranges.add_control_edge(fork, self.expr_count + 1);
+                self.drop_ranges.add_control_edge(fork, self.expr_index + 1);
                 if let Some(if_false) = if_false {
                     self.visit_expr(if_false);
                 }
 
-                self.drop_ranges.add_control_edge(true_end, self.expr_count + 1);
+                self.drop_ranges.add_control_edge(true_end, self.expr_index + 1);
             }
             ExprKind::Assign(lhs, rhs, _) => {
                 self.visit_expr(lhs);
@@ -122,18 +126,18 @@ impl<'tcx> Visitor<'tcx> for DropRangeVisitor<'tcx> {
                 reinit = Some(lhs);
             }
             ExprKind::Loop(body, ..) => {
-                let loop_begin = self.expr_count + 1;
+                let loop_begin = self.expr_index + 1;
                 self.visit_block(body);
-                self.drop_ranges.add_control_edge(self.expr_count, loop_begin);
+                self.drop_ranges.add_control_edge(self.expr_index, loop_begin);
             }
             ExprKind::Match(scrutinee, arms, ..) => {
                 self.visit_expr(scrutinee);
 
-                let fork = self.expr_count;
+                let fork = self.expr_index;
                 let arm_end_ids = arms
                     .iter()
                     .map(|hir::Arm { pat, body, guard, .. }| {
-                        self.drop_ranges.add_control_edge(fork, self.expr_count + 1);
+                        self.drop_ranges.add_control_edge(fork, self.expr_index + 1);
                         self.visit_pat(pat);
                         match guard {
                             Some(Guard::If(expr)) => self.visit_expr(expr),
@@ -144,23 +148,23 @@ impl<'tcx> Visitor<'tcx> for DropRangeVisitor<'tcx> {
                             None => (),
                         }
                         self.visit_expr(body);
-                        self.expr_count
+                        self.expr_index
                     })
                     .collect::<Vec<_>>();
                 arm_end_ids.into_iter().for_each(|arm_end| {
-                    self.drop_ranges.add_control_edge(arm_end, self.expr_count + 1)
+                    self.drop_ranges.add_control_edge(arm_end, self.expr_index + 1)
                 });
             }
             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_count, target);
+                self.drop_ranges.add_control_edge_hir_id(self.expr_index, target);
             }
 
             _ => intravisit::walk_expr(self, expr),
         }
 
-        self.expr_count += 1;
-        self.drop_ranges.add_node_mapping(expr.hir_id, self.expr_count);
+        self.expr_index = self.expr_index + 1;
+        self.drop_ranges.add_node_mapping(expr.hir_id, self.expr_index);
         self.consume_expr(expr);
         if let Some(expr) = reinit {
             self.reinit_expr(expr);
@@ -171,11 +175,11 @@ impl<'tcx> Visitor<'tcx> for DropRangeVisitor<'tcx> {
         intravisit::walk_pat(self, pat);
 
         // Increment expr_count here to match what InteriorVisitor expects.
-        self.expr_count += 1;
+        self.expr_index = self.expr_index + 1;
     }
 }
 
-impl DropRanges {
+impl DropRangesBuilder {
     fn new(hir_ids: impl Iterator<Item = HirId>, hir: Map<'_>, num_exprs: usize) -> Self {
         let mut hir_id_map = HirIdMap::<HirIdIndex>::default();
         let mut next = <_>::from(0u32);
@@ -204,7 +208,7 @@ impl DropRanges {
     /// Adds an entry in the mapping from HirIds to PostOrderIds
     ///
     /// Needed so that `add_control_edge_hir_id` can work.
-    fn add_node_mapping(&mut self, hir_id: HirId, post_order_id: usize) {
+    fn add_node_mapping(&mut self, hir_id: HirId, post_order_id: PostOrderId) {
         self.post_order_map.insert(hir_id, post_order_id);
     }
 
@@ -212,16 +216,16 @@ impl DropRanges {
     ///
     /// This can be used for branches where we do not know the PostOrderId of the target yet,
     /// such as when handling `break` or `continue`.
-    fn add_control_edge_hir_id(&mut self, from: usize, to: HirId) {
+    fn add_control_edge_hir_id(&mut self, from: PostOrderId, to: HirId) {
         self.deferred_edges.push((from, to));
     }
 
-    fn drop_at(&mut self, value: HirId, location: usize) {
+    fn drop_at(&mut self, value: HirId, location: PostOrderId) {
         let value = self.hidx(value);
         self.node_mut(location.into()).drops.push(value);
     }
 
-    fn reinit_at(&mut self, value: HirId, location: usize) {
+    fn reinit_at(&mut self, value: HirId, location: PostOrderId) {
         let value = match self.hir_id_map.get(&value) {
             Some(value) => *value,
             // If there's no value, this is never consumed and therefore is never dropped. We can
@@ -230,4 +234,18 @@ impl DropRanges {
         };
         self.node_mut(location.into()).reinits.push(value);
     }
+
+    /// Looks up PostOrderId for any control edges added by HirId and adds a proper edge for them.
+    ///
+    /// 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) {
+        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);
+            self.add_control_edge(from, to)
+        });
+    }
 }
diff --git a/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/cfg_propagate.rs b/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/cfg_propagate.rs
index 74ce762864e..22f7484abf3 100644
--- a/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/cfg_propagate.rs
+++ b/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/cfg_propagate.rs
@@ -1,12 +1,10 @@
-use super::{DropRanges, PostOrderId};
+use super::{DropRangesBuilder, PostOrderId};
 use rustc_index::{bit_set::BitSet, vec::IndexVec};
 use std::collections::BTreeMap;
-use std::mem::swap;
 
-impl DropRanges {
+impl DropRangesBuilder {
     pub fn propagate_to_fixpoint(&mut self) {
         trace!("before fixpoint: {:#?}", self);
-        self.process_deferred_edges();
         let preds = self.compute_predecessors();
 
         trace!("predecessors: {:#?}", preds.iter_enumerated().collect::<BTreeMap<_, _>>());
@@ -53,6 +51,11 @@ impl DropRanges {
     fn compute_predecessors(&self) -> IndexVec<PostOrderId, Vec<PostOrderId>> {
         let mut preds = IndexVec::from_fn_n(|_| vec![], self.nodes.len());
         for (id, node) in self.nodes.iter_enumerated() {
+            // If the node has no explicit successors, we assume that control
+            // will from this node into the next one.
+            //
+            // If there are successors listed, then we assume that all
+            // possible successors are given and we do not include the default.
             if node.successors.len() == 0 && id.index() != self.nodes.len() - 1 {
                 preds[id + 1].push(id);
             } else {
@@ -63,18 +66,4 @@ impl DropRanges {
         }
         preds
     }
-
-    /// Looks up PostOrderId for any control edges added by HirId and adds a proper edge for them.
-    ///
-    /// 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) {
-        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);
-            self.add_control_edge(from, to)
-        });
-    }
 }
diff --git a/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/cfg_visualize.rs b/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/cfg_visualize.rs
index ebbbec1c472..b87b3dd9a5f 100644
--- a/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/cfg_visualize.rs
+++ b/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/cfg_visualize.rs
@@ -3,9 +3,9 @@
 
 use rustc_graphviz as dot;
 
-use super::{DropRanges, PostOrderId};
+use super::{DropRangesBuilder, PostOrderId};
 
-impl<'a> dot::GraphWalk<'a> for DropRanges {
+impl<'a> dot::GraphWalk<'a> for DropRangesBuilder {
     type Node = PostOrderId;
 
     type Edge = (PostOrderId, PostOrderId);
@@ -36,7 +36,7 @@ impl<'a> dot::GraphWalk<'a> for DropRanges {
     }
 }
 
-impl<'a> dot::Labeller<'a> for DropRanges {
+impl<'a> dot::Labeller<'a> for DropRangesBuilder {
     type Node = PostOrderId;
 
     type Edge = (PostOrderId, PostOrderId);
@@ -56,7 +56,7 @@ impl<'a> dot::Labeller<'a> for DropRanges {
                 n,
                 self.post_order_map
                     .iter()
-                    .find(|(_hir_id, &post_order_id)| post_order_id == n.index())
+                    .find(|(_hir_id, &post_order_id)| post_order_id == *n)
                     .map_or("<unknown>".into(), |(hir_id, _)| format!(
                         "{}",
                         hir_id.local_id.index()