Make drop-use fact collection simpler for `polonius`

This shunts all the complexity of siphoning off the drop-use facts
into `LivenessResults::add_extra_drop_facts()`, which may or may
not be a good approach.

4 files changed, 34 insertions, 41 deletions

diff --git a/compiler/rustc_borrowck/src/type_check/liveness/mod.rs b/compiler/rustc_borrowck/src/type_check/liveness/mod.rs
index 38ec9f7678e..8b863efad6c 100644
--- a/compiler/rustc_borrowck/src/type_check/liveness/mod.rs
+++ b/compiler/rustc_borrowck/src/type_check/liveness/mod.rs
@@ -14,7 +14,6 @@ use std::rc::Rc;
 use crate::{
     constraints::OutlivesConstraintSet,
     facts::{AllFacts, AllFactsExt},
-    location::LocationTable,
     region_infer::values::LivenessValues,
     universal_regions::UniversalRegions,
 };
@@ -39,7 +38,6 @@ pub(super) fn generate<'mir, 'tcx>(
     elements: &Rc<DenseLocationMap>,
     flow_inits: &mut ResultsCursor<'mir, 'tcx, MaybeInitializedPlaces<'mir, 'tcx>>,
     move_data: &MoveData<'tcx>,
-    location_table: &LocationTable,
     use_polonius: bool,
 ) {
     debug!("liveness::generate");
@@ -53,11 +51,9 @@ pub(super) fn generate<'mir, 'tcx>(
         compute_relevant_live_locals(typeck.tcx(), &free_regions, body);
     let facts_enabled = use_polonius || AllFacts::enabled(typeck.tcx());
 
-    let polonius_drop_used = facts_enabled.then(|| {
-        let mut drop_used = Vec::new();
-        polonius::populate_access_facts(typeck, body, location_table, move_data, &mut drop_used);
-        drop_used
-    });
+    if facts_enabled {
+        polonius::populate_access_facts(typeck, body, move_data);
+    };
 
     trace::trace(
         typeck,
@@ -67,7 +63,6 @@ pub(super) fn generate<'mir, 'tcx>(
         move_data,
         relevant_live_locals,
         boring_locals,
-        polonius_drop_used,
     );
 
     // Mark regions that should be live where they appear within rvalues or within a call: like
diff --git a/compiler/rustc_borrowck/src/type_check/liveness/polonius.rs b/compiler/rustc_borrowck/src/type_check/liveness/polonius.rs
index 808df1d66cb..d8f03a07a63 100644
--- a/compiler/rustc_borrowck/src/type_check/liveness/polonius.rs
+++ b/compiler/rustc_borrowck/src/type_check/liveness/polonius.rs
@@ -14,7 +14,7 @@ struct UseFactsExtractor<'me, 'tcx> {
     var_defined_at: &'me mut VarPointRelation,
     var_used_at: &'me mut VarPointRelation,
     location_table: &'me LocationTable,
-    var_dropped_at: &'me mut Vec<(Local, Location)>,
+    var_dropped_at: &'me mut VarPointRelation,
     move_data: &'me MoveData<'tcx>,
     path_accessed_at_base: &'me mut PathPointRelation,
 }
@@ -37,7 +37,7 @@ impl<'tcx> UseFactsExtractor<'_, 'tcx> {
 
     fn insert_drop_use(&mut self, local: Local, location: Location) {
         debug!("UseFactsExtractor::insert_drop_use()");
-        self.var_dropped_at.push((local, location));
+        self.var_dropped_at.push((local, self.location_to_index(location)));
     }
 
     fn insert_path_access(&mut self, path: MovePathIndex, location: Location) {
@@ -85,33 +85,22 @@ impl<'a, 'tcx> Visitor<'tcx> for UseFactsExtractor<'a, 'tcx> {
 pub(super) fn populate_access_facts<'a, 'tcx>(
     typeck: &mut TypeChecker<'a, 'tcx>,
     body: &Body<'tcx>,
-    location_table: &LocationTable,
     move_data: &MoveData<'tcx>,
-    // FIXME: this is an inelegant way of squirreling away a
-    // copy of `var_dropped_at` in the original `Location` format
-    // for later use in `trace::trace()`, which updates some liveness-
-    // internal data based on what Polonius saw.
-    // Ideally, that part would access the Polonius facts directly, and this
-    // would be regular facts gathering.
-    dropped_at: &mut Vec<(Local, Location)>,
 ) {
     debug!("populate_access_facts()");
+    let location_table = typeck.borrowck_context.location_table;
 
     if let Some(facts) = typeck.borrowck_context.all_facts.as_mut() {
         let mut extractor = UseFactsExtractor {
             var_defined_at: &mut facts.var_defined_at,
             var_used_at: &mut facts.var_used_at,
-            var_dropped_at: dropped_at,
+            var_dropped_at: &mut facts.var_dropped_at,
             path_accessed_at_base: &mut facts.path_accessed_at_base,
             location_table,
             move_data,
         };
         extractor.visit_body(body);
 
-        facts.var_dropped_at.extend(
-            dropped_at.iter().map(|&(local, location)| (local, location_table.mid_index(location))),
-        );
-
         for (local, local_decl) in body.local_decls.iter_enumerated() {
             debug!(
                 "add use_of_var_derefs_origin facts - local={:?}, type={:?}",
diff --git a/compiler/rustc_borrowck/src/type_check/liveness/trace.rs b/compiler/rustc_borrowck/src/type_check/liveness/trace.rs
index 6cc0e67c0f8..92b70a09a88 100644
--- a/compiler/rustc_borrowck/src/type_check/liveness/trace.rs
+++ b/compiler/rustc_borrowck/src/type_check/liveness/trace.rs
@@ -16,6 +16,7 @@ use rustc_mir_dataflow::impls::MaybeInitializedPlaces;
 use rustc_mir_dataflow::move_paths::{HasMoveData, MoveData, MovePathIndex};
 use rustc_mir_dataflow::ResultsCursor;
 
+use crate::location::RichLocation;
 use crate::{
     region_infer::values::{self, LiveLoans},
     type_check::liveness::local_use_map::LocalUseMap,
@@ -46,7 +47,6 @@ pub(super) fn trace<'mir, 'tcx>(
     move_data: &MoveData<'tcx>,
     relevant_live_locals: Vec<Local>,
     boring_locals: Vec<Local>,
-    polonius_drop_used: Option<Vec<(Local, Location)>>,
 ) {
     let local_use_map = &LocalUseMap::build(&relevant_live_locals, elements, body);
 
@@ -81,6 +81,8 @@ pub(super) fn trace<'mir, 'tcx>(
         borrowck_context.constraints.liveness_constraints.loans = Some(live_loans);
     };
 
+    let polonius_facts_gathered = typeck.borrowck_context.all_facts.is_some();
+
     let cx = LivenessContext {
         typeck,
         body,
@@ -93,8 +95,8 @@ pub(super) fn trace<'mir, 'tcx>(
 
     let mut results = LivenessResults::new(cx);
 
-    if let Some(drop_used) = polonius_drop_used {
-        results.add_extra_drop_facts(drop_used, relevant_live_locals.iter().copied().collect())
+    if polonius_facts_gathered {
+        results.add_extra_drop_facts(relevant_live_locals.iter().copied().collect());
     }
 
     results.compute_for_all_locals(relevant_live_locals);
@@ -218,17 +220,32 @@ impl<'me, 'typeck, 'flow, 'tcx> LivenessResults<'me, 'typeck, 'flow, 'tcx> {
     ///
     /// Add facts for all locals with free regions, since regions may outlive
     /// the function body only at certain nodes in the CFG.
-    fn add_extra_drop_facts(
-        &mut self,
-        drop_used: Vec<(Local, Location)>,
-        relevant_live_locals: FxIndexSet<Local>,
-    ) {
+    fn add_extra_drop_facts(&mut self, relevant_live_locals: FxIndexSet<Local>) {
+        let drop_used = self
+            .cx
+            .typeck
+            .borrowck_context
+            .all_facts
+            .as_ref()
+            .map(|facts| facts.var_dropped_at.clone())
+            .into_iter()
+            .flatten();
         let locations = IntervalSet::new(self.cx.elements.num_points());
 
-        for (local, location) in drop_used {
+        for (local, location_index) in drop_used {
             if !relevant_live_locals.contains(&local) {
                 let local_ty = self.cx.body.local_decls[local].ty;
                 if local_ty.has_free_regions() {
+                    let location = match self
+                        .cx
+                        .typeck
+                        .borrowck_context
+                        .location_table
+                        .to_location(location_index)
+                    {
+                        RichLocation::Start(l) => l,
+                        RichLocation::Mid(l) => l,
+                    };
                     self.cx.add_drop_live_facts_for(local, local_ty, &[location], &locations);
                 }
             }
diff --git a/compiler/rustc_borrowck/src/type_check/mod.rs b/compiler/rustc_borrowck/src/type_check/mod.rs
index 4e46a0c62c7..a05fa967af0 100644
--- a/compiler/rustc_borrowck/src/type_check/mod.rs
+++ b/compiler/rustc_borrowck/src/type_check/mod.rs
@@ -188,15 +188,7 @@ pub(crate) fn type_check<'mir, 'tcx>(
     checker.equate_inputs_and_outputs(body, universal_regions, &normalized_inputs_and_output);
     checker.check_signature_annotation(body);
 
-    liveness::generate(
-        &mut checker,
-        body,
-        elements,
-        flow_inits,
-        move_data,
-        location_table,
-        use_polonius,
-    );
+    liveness::generate(&mut checker, body, elements, flow_inits, move_data, use_polonius);
 
     translate_outlives_facts(&mut checker);
     let opaque_type_values = infcx.take_opaque_types();