From c48756cdbfb1725251cbfa6fe760b2cb4e47b2d9 Mon Sep 17 00:00:00 2001
From: Camille GILLOT <gillot.camille@gmail.com>
Date: Sat, 28 Jan 2023 11:23:18 +0000
Subject: [PATCH] Limit creation of tracked place directly.

---
 .../rustc_mir_dataflow/src/value_analysis.rs  | 22 +++++++++++++++----
 .../src/dataflow_const_prop.rs                | 11 ++++------
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/compiler/rustc_mir_dataflow/src/value_analysis.rs b/compiler/rustc_mir_dataflow/src/value_analysis.rs
index 03b6c182062..f587f17f12a 100644
--- a/compiler/rustc_mir_dataflow/src/value_analysis.rs
+++ b/compiler/rustc_mir_dataflow/src/value_analysis.rs
@@ -599,10 +599,11 @@ impl Map {
         tcx: TyCtxt<'tcx>,
         body: &Body<'tcx>,
         filter: impl FnMut(Ty<'tcx>) -> bool,
+        place_limit: Option<usize>,
     ) -> Self {
         let mut map = Self::new();
         let exclude = excluded_locals(body);
-        map.register_with_filter(tcx, body, filter, exclude);
+        map.register_with_filter(tcx, body, filter, exclude, place_limit);
         debug!("registered {} places ({} nodes in total)", map.value_count, map.places.len());
         map
     }
@@ -614,12 +615,20 @@ impl Map {
         body: &Body<'tcx>,
         mut filter: impl FnMut(Ty<'tcx>) -> bool,
         exclude: BitSet<Local>,
+        place_limit: Option<usize>,
     ) {
         // We use this vector as stack, pushing and popping projections.
         let mut projection = Vec::new();
         for (local, decl) in body.local_decls.iter_enumerated() {
             if !exclude.contains(local) {
-                self.register_with_filter_rec(tcx, local, &mut projection, decl.ty, &mut filter);
+                self.register_with_filter_rec(
+                    tcx,
+                    local,
+                    &mut projection,
+                    decl.ty,
+                    &mut filter,
+                    place_limit,
+                );
             }
         }
     }
@@ -634,7 +643,12 @@ impl Map {
         projection: &mut Vec<PlaceElem<'tcx>>,
         ty: Ty<'tcx>,
         filter: &mut impl FnMut(Ty<'tcx>) -> bool,
+        place_limit: Option<usize>,
     ) {
+        if let Some(place_limit) = place_limit && self.value_count >= place_limit {
+            return
+        }
+
         // We know that the projection only contains trackable elements.
         let place = self.make_place(local, projection).unwrap();
 
@@ -672,13 +686,13 @@ impl Map {
                 projection.push(PlaceElem::Downcast(None, variant));
                 let _ = self.make_place(local, projection);
                 projection.push(PlaceElem::Field(field, ty));
-                self.register_with_filter_rec(tcx, local, projection, ty, filter);
+                self.register_with_filter_rec(tcx, local, projection, ty, filter, place_limit);
                 projection.pop();
                 projection.pop();
                 return;
             }
             projection.push(PlaceElem::Field(field, ty));
-            self.register_with_filter_rec(tcx, local, projection, ty, filter);
+            self.register_with_filter_rec(tcx, local, projection, ty, filter, place_limit);
             projection.pop();
         });
     }
diff --git a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
index f10f208f5de..bfb1eb8b5fb 100644
--- a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
+++ b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
@@ -37,9 +37,6 @@ impl<'tcx> MirPass<'tcx> for DataflowConstProp {
             return;
         }
 
-        // Decide which places to track during the analysis.
-        let map = Map::from_filter(tcx, body, Ty::is_scalar);
-
         // We want to have a somewhat linear runtime w.r.t. the number of statements/terminators.
         // Let's call this number `n`. Dataflow analysis has `O(h*n)` transfer function
         // applications, where `h` is the height of the lattice. Because the height of our lattice
@@ -48,10 +45,10 @@ impl<'tcx> MirPass<'tcx> for DataflowConstProp {
         // `O(num_nodes * tracked_places * n)` in terms of time complexity. Since the number of
         // map nodes is strongly correlated to the number of tracked places, this becomes more or
         // less `O(n)` if we place a constant limit on the number of tracked places.
-        if tcx.sess.mir_opt_level() < 4 && map.tracked_places() > PLACE_LIMIT {
-            debug!("aborted dataflow const prop due to too many tracked places");
-            return;
-        }
+        let place_limit = if tcx.sess.mir_opt_level() < 4 { Some(PLACE_LIMIT) } else { None };
+
+        // Decide which places to track during the analysis.
+        let map = Map::from_filter(tcx, body, Ty::is_scalar, place_limit);
 
         // Perform the actual dataflow analysis.
         let analysis = ConstAnalysis::new(tcx, body, map);