drop_ranges: Add TrackedValue enum

This makes it clearer what values we are tracking and why.

3 files changed, 106 insertions, 67 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 cf463d0aeae..681cd7cf935 100644
--- a/compiler/rustc_typeck/src/check/generator_interior/drop_ranges.rs
+++ b/compiler/rustc_typeck/src/check/generator_interior/drop_ranges.rs
@@ -17,9 +17,12 @@ use self::record_consumed_borrow::find_consumed_and_borrowed;
 use crate::check::FnCtxt;
 use hir::def_id::DefId;
 use hir::{Body, HirId, HirIdMap, Node};
+use rustc_data_structures::fx::FxHashMap;
 use rustc_hir as hir;
 use rustc_index::bit_set::BitSet;
 use rustc_index::vec::IndexVec;
+use rustc_middle::hir::place::{PlaceBase, PlaceWithHirId};
+use rustc_middle::ty;
 use std::collections::BTreeMap;
 use std::fmt::Debug;
 
@@ -47,13 +50,17 @@ pub fn compute_drop_ranges<'a, 'tcx>(
 
     drop_ranges.propagate_to_fixpoint();
 
-    DropRanges { hir_id_map: drop_ranges.hir_id_map, nodes: drop_ranges.nodes }
+    DropRanges { tracked_value_map: drop_ranges.tracked_value_map, nodes: drop_ranges.nodes }
 }
 
 /// Applies `f` to consumable portion of a HIR node.
 ///
 /// The `node` parameter should be the result of calling `Map::find(place)`.
-fn for_each_consumable(place: HirId, node: Option<Node<'_>>, mut f: impl FnMut(HirId)) {
+fn for_each_consumable(
+    place: TrackedValue,
+    node: Option<Node<'_>>,
+    mut f: impl FnMut(TrackedValue),
+) {
     f(place);
     if let Some(Node::Expr(expr)) = node {
         match expr.kind {
@@ -61,7 +68,7 @@ fn for_each_consumable(place: HirId, node: Option<Node<'_>>, mut f: impl FnMut(H
                 _,
                 hir::Path { res: hir::def::Res::Local(hir_id), .. },
             )) => {
-                f(*hir_id);
+                f(TrackedValue::Variable(*hir_id));
             }
             _ => (),
         }
@@ -75,22 +82,60 @@ rustc_index::newtype_index! {
 }
 
 rustc_index::newtype_index! {
-    pub struct HirIdIndex {
+    pub struct TrackedValueIndex {
         DEBUG_FORMAT = "hidx({})",
     }
 }
 
+/// Identifies a value whose drop state we need to track.
+#[derive(PartialEq, Eq, Hash, Debug, Clone, Copy)]
+enum TrackedValue {
+    /// Represents a named variable, such as a let binding, parameter, or upvar.
+    ///
+    /// The HirId points to the variable's definition site.
+    Variable(HirId),
+    /// A value produced as a result of an expression.
+    ///
+    /// The HirId points to the expression that returns this value.
+    Temporary(HirId),
+}
+
+impl TrackedValue {
+    fn hir_id(&self) -> HirId {
+        match self {
+            TrackedValue::Variable(hir_id) | TrackedValue::Temporary(hir_id) => *hir_id,
+        }
+    }
+}
+
+impl From<&PlaceWithHirId<'_>> for TrackedValue {
+    fn from(place_with_id: &PlaceWithHirId<'_>) -> Self {
+        match place_with_id.place.base {
+            PlaceBase::Rvalue | PlaceBase::StaticItem => {
+                TrackedValue::Temporary(place_with_id.hir_id)
+            }
+            PlaceBase::Local(hir_id)
+            | PlaceBase::Upvar(ty::UpvarId { var_path: ty::UpvarPath { hir_id }, .. }) => {
+                TrackedValue::Variable(hir_id)
+            }
+        }
+    }
+}
+
 pub struct DropRanges {
-    hir_id_map: HirIdMap<HirIdIndex>,
+    tracked_value_map: FxHashMap<TrackedValue, TrackedValueIndex>,
     nodes: IndexVec<PostOrderId, NodeInfo>,
 }
 
 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))
+        self.tracked_value_map
+            .get(&TrackedValue::Temporary(hir_id))
+            .or(self.tracked_value_map.get(&TrackedValue::Variable(hir_id)))
+            .cloned()
+            .map_or(false, |tracked_value_id| {
+                self.expect_node(location.into()).drop_state.contains(tracked_value_id)
+            })
     }
 
     /// Returns a reference to the NodeInfo for a node, panicking if it does not exist
@@ -118,7 +163,7 @@ struct DropRangesBuilder {
     /// (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>,
+    tracked_value_map: FxHashMap<TrackedValue, TrackedValueIndex>,
 
     /// 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.
@@ -138,7 +183,7 @@ struct DropRangesBuilder {
 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)
+            .field("hir_id_map", &self.tracked_value_map)
             .field("post_order_maps", &self.post_order_map)
             .field("nodes", &self.nodes.iter_enumerated().collect::<BTreeMap<_, _>>())
             .finish()
@@ -154,7 +199,7 @@ impl Debug for DropRangesBuilder {
 impl DropRangesBuilder {
     /// Returns the number of values (hir_ids) that are tracked
     fn num_values(&self) -> usize {
-        self.hir_id_map.len()
+        self.tracked_value_map.len()
     }
 
     fn node_mut(&mut self, id: PostOrderId) -> &mut NodeInfo {
@@ -177,13 +222,13 @@ struct NodeInfo {
     successors: Vec<PostOrderId>,
 
     /// List of hir_ids that are dropped by this node.
-    drops: Vec<HirIdIndex>,
+    drops: Vec<TrackedValueIndex>,
 
     /// List of hir_ids that are reinitialized by this node.
-    reinits: Vec<HirIdIndex>,
+    reinits: Vec<TrackedValueIndex>,
 
     /// Set of values that are definitely dropped at this point.
-    drop_state: BitSet<HirIdIndex>,
+    drop_state: BitSet<TrackedValueIndex>,
 }
 
 impl NodeInfo {
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 1591b144dc6..dfe8ed54b21 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,11 +1,12 @@
 use super::{
     for_each_consumable, record_consumed_borrow::ConsumedAndBorrowedPlaces, DropRangesBuilder,
-    HirIdIndex, NodeInfo, PostOrderId,
+    NodeInfo, PostOrderId, TrackedValue, TrackedValueIndex,
 };
 use hir::{
     intravisit::{self, NestedVisitorMap, Visitor},
-    Body, Expr, ExprKind, Guard, HirId, HirIdMap,
+    Body, Expr, ExprKind, Guard, HirId,
 };
+use rustc_data_structures::fx::FxHashMap;
 use rustc_hir as hir;
 use rustc_index::vec::IndexVec;
 use rustc_middle::{
@@ -61,20 +62,20 @@ impl<'a, 'tcx> DropRangeVisitor<'a, 'tcx> {
     ) -> Self {
         debug!("consumed_places: {:?}", places.consumed);
         let drop_ranges = DropRangesBuilder::new(
-            places.consumed.iter().flat_map(|(_, places)| places.iter().copied()),
+            places.consumed.iter().flat_map(|(_, places)| places.iter().cloned()),
             hir,
             num_exprs,
         );
         Self { hir, places, drop_ranges, expr_index: PostOrderId::from_u32(0), typeck_results, tcx }
     }
 
-    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);
+    fn record_drop(&mut self, value: TrackedValue) {
+        if self.places.borrowed.contains(&value) {
+            debug!("not marking {:?} as dropped because it is borrowed at some point", value);
         } else {
-            debug!("marking {:?} as dropped at {:?}", hir_id, self.expr_index);
+            debug!("marking {:?} as dropped at {:?}", value, self.expr_index);
             let count = self.expr_index;
-            self.drop_ranges.drop_at(hir_id, count);
+            self.drop_ranges.drop_at(value, count);
         }
     }
 
@@ -88,7 +89,9 @@ impl<'a, 'tcx> DropRangeVisitor<'a, 'tcx> {
             .get(&expr.hir_id)
             .map_or(vec![], |places| places.iter().cloned().collect());
         for place in places {
-            for_each_consumable(place, self.hir.find(place), |hir_id| self.record_drop(hir_id));
+            for_each_consumable(place, self.hir.find(place.hir_id()), |value| {
+                self.record_drop(value)
+            });
         }
     }
 
@@ -100,7 +103,7 @@ impl<'a, 'tcx> DropRangeVisitor<'a, 'tcx> {
         {
             let location = self.expr_index;
             debug!("reinitializing {:?} at {:?}", hir_id, location);
-            self.drop_ranges.reinit_at(*hir_id, location);
+            self.drop_ranges.reinit_at(TrackedValue::Variable(*hir_id), location);
         } else {
             debug!("reinitializing {:?} is not supported", expr);
         }
@@ -264,36 +267,40 @@ impl<'a, 'tcx> Visitor<'tcx> for DropRangeVisitor<'a, 'tcx> {
 }
 
 impl DropRangesBuilder {
-    fn new(hir_ids: impl Iterator<Item = HirId>, hir: Map<'_>, num_exprs: usize) -> Self {
-        let mut hir_id_map = HirIdMap::<HirIdIndex>::default();
+    fn new(
+        tracked_values: impl Iterator<Item = TrackedValue>,
+        hir: Map<'_>,
+        num_exprs: usize,
+    ) -> Self {
+        let mut tracked_value_map = FxHashMap::<_, TrackedValueIndex>::default();
         let mut next = <_>::from(0u32);
-        for hir_id in hir_ids {
-            for_each_consumable(hir_id, hir.find(hir_id), |hir_id| {
-                if !hir_id_map.contains_key(&hir_id) {
-                    hir_id_map.insert(hir_id, next);
-                    next = <_>::from(next.index() + 1);
+        for value in tracked_values {
+            for_each_consumable(value, hir.find(value.hir_id()), |value| {
+                if !tracked_value_map.contains_key(&value) {
+                    tracked_value_map.insert(value, next);
+                    next = next + 1;
                 }
             });
         }
-        debug!("hir_id_map: {:?}", hir_id_map);
-        let num_values = hir_id_map.len();
+        debug!("hir_id_map: {:?}", tracked_value_map);
+        let num_values = tracked_value_map.len();
         Self {
-            hir_id_map,
+            tracked_value_map,
             nodes: IndexVec::from_fn_n(|_| NodeInfo::new(num_values), num_exprs + 1),
             deferred_edges: <_>::default(),
             post_order_map: <_>::default(),
         }
     }
 
-    fn hidx(&self, hir_id: HirId) -> HirIdIndex {
-        *self.hir_id_map.get(&hir_id).unwrap()
+    fn tracked_value_index(&self, tracked_value: TrackedValue) -> TrackedValueIndex {
+        *self.tracked_value_map.get(&tracked_value).unwrap()
     }
 
     /// 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: PostOrderId) {
-        self.post_order_map.insert(hir_id, post_order_id);
+    fn add_node_mapping(&mut self, node_hir_id: HirId, post_order_id: PostOrderId) {
+        self.post_order_map.insert(node_hir_id, post_order_id);
     }
 
     /// Like add_control_edge, but uses a hir_id as the target.
@@ -304,13 +311,13 @@ impl DropRangesBuilder {
         self.deferred_edges.push((from, to));
     }
 
-    fn drop_at(&mut self, value: HirId, location: PostOrderId) {
-        let value = self.hidx(value);
+    fn drop_at(&mut self, value: TrackedValue, location: PostOrderId) {
+        let value = self.tracked_value_index(value);
         self.node_mut(location.into()).drops.push(value);
     }
 
-    fn reinit_at(&mut self, value: HirId, location: PostOrderId) {
-        let value = match self.hir_id_map.get(&value) {
+    fn reinit_at(&mut self, value: TrackedValue, location: PostOrderId) {
+        let value = match self.tracked_value_map.get(&value) {
             Some(value) => *value,
             // If there's no value, this is never consumed and therefore is never dropped. We can
             // ignore this.
diff --git a/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/record_consumed_borrow.rs b/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/record_consumed_borrow.rs
index 36e843e7fd1..2548b608092 100644
--- a/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/record_consumed_borrow.rs
+++ b/compiler/rustc_typeck/src/check/generator_interior/drop_ranges/record_consumed_borrow.rs
@@ -1,16 +1,14 @@
+use super::TrackedValue;
 use crate::{
     check::FnCtxt,
     expr_use_visitor::{self, ExprUseVisitor},
 };
-use hir::{def_id::DefId, Body, HirId, HirIdMap, HirIdSet};
+use hir::{def_id::DefId, Body, HirId, HirIdMap};
+use rustc_data_structures::stable_set::FxHashSet;
 use rustc_hir as hir;
-use rustc_middle::hir::{
-    map::Map,
-    place::{Place, PlaceBase},
-};
-use rustc_middle::ty;
+use rustc_middle::hir::map::Map;
 
-pub fn find_consumed_and_borrowed<'a, 'tcx>(
+pub(super) fn find_consumed_and_borrowed<'a, 'tcx>(
     fcx: &'a FnCtxt<'a, 'tcx>,
     def_id: DefId,
     body: &'tcx Body<'tcx>,
@@ -20,7 +18,7 @@ pub fn find_consumed_and_borrowed<'a, 'tcx>(
     expr_use_visitor.places
 }
 
-pub struct ConsumedAndBorrowedPlaces {
+pub(super) struct ConsumedAndBorrowedPlaces {
     /// Records the variables/expressions that are dropped by a given expression.
     ///
     /// The key is the hir-id of the expression, and the value is a set or hir-ids for variables
@@ -28,9 +26,9 @@ pub struct ConsumedAndBorrowedPlaces {
     ///
     /// Note that this set excludes "partial drops" -- for example, a statement like `drop(x.y)` is
     /// not considered a drop of `x`, although it would be a drop of `x.y`.
-    pub consumed: HirIdMap<HirIdSet>,
+    pub(super) consumed: HirIdMap<FxHashSet<TrackedValue>>,
     /// A set of hir-ids of values or variables that are borrowed at some point within the body.
-    pub borrowed: HirIdSet,
+    pub(super) borrowed: FxHashSet<TrackedValue>,
 }
 
 /// Works with ExprUseVisitor to find interesting values for the drop range analysis.
@@ -65,7 +63,7 @@ impl<'tcx> ExprUseDelegate<'tcx> {
         .consume_body(body);
     }
 
-    fn mark_consumed(&mut self, consumer: HirId, target: HirId) {
+    fn mark_consumed(&mut self, consumer: HirId, target: TrackedValue) {
         if !self.places.consumed.contains_key(&consumer) {
             self.places.consumed.insert(consumer, <_>::default());
         }
@@ -87,8 +85,7 @@ impl<'tcx> expr_use_visitor::Delegate<'tcx> for ExprUseDelegate<'tcx> {
             "consume {:?}; diag_expr_id={:?}, using parent {:?}",
             place_with_id, diag_expr_id, parent
         );
-        self.mark_consumed(parent, place_with_id.hir_id);
-        place_hir_id(&place_with_id.place).map(|place| self.mark_consumed(parent, place));
+        self.mark_consumed(parent, place_with_id.into());
     }
 
     fn borrow(
@@ -97,7 +94,7 @@ impl<'tcx> expr_use_visitor::Delegate<'tcx> for ExprUseDelegate<'tcx> {
         _diag_expr_id: HirId,
         _bk: rustc_middle::ty::BorrowKind,
     ) {
-        place_hir_id(&place_with_id.place).map(|place| self.places.borrowed.insert(place));
+        self.places.borrowed.insert(place_with_id.into());
     }
 
     fn mutate(
@@ -115,13 +112,3 @@ impl<'tcx> expr_use_visitor::Delegate<'tcx> for ExprUseDelegate<'tcx> {
     ) {
     }
 }
-
-/// Gives the hir_id associated with a place if one exists. This is the hir_id that we want to
-/// track for a value in the drop range analysis.
-fn place_hir_id(place: &Place<'_>) -> Option<HirId> {
-    match place.base {
-        PlaceBase::Rvalue | PlaceBase::StaticItem => None,
-        PlaceBase::Local(hir_id)
-        | PlaceBase::Upvar(ty::UpvarId { var_path: ty::UpvarPath { hir_id }, .. }) => Some(hir_id),
-    }
-}