Auto merge of rust-lang#127087 - cjgillot:small-map, r=<try>

Only track mentionned places for jump threading

This PR aims to reduce the state space size in jump threading and dataflow const-prop opts.

The current implementation walks the types of all locals, and creates a place for each possible projection. This can easily lead to a large number of places and tracked values, most being useless to the actual pass.

With this PR, we instead collect places that appear syntactically in the MIR (first commit). However, this is not sufficient (second commit), and we miss places that we could track in aggregate assignments. The third commit tracks such assignments to mirror place projections, see the inline comment.

This is complementary to rust-lang#127036

r? `@oli-obk`

Author: bors

compiler/rustc_mir_dataflow/src/value_analysis.rs

@@ -66,7 +66,7 @@ impl<'tcx> MirPass<'tcx> for DataflowConstProp {
 }
 
 struct ConstAnalysis<'a, 'tcx> {
-    map: Map,
+    map: Map<'tcx>,
     tcx: TyCtxt<'tcx>,
     local_decls: &'a LocalDecls<'tcx>,
     ecx: InterpCx<'tcx, DummyMachine>,
@@ -78,7 +78,7 @@ impl<'tcx> ValueAnalysis<'tcx> for ConstAnalysis<'_, 'tcx> {
 
     const NAME: &'static str = "ConstAnalysis";
 
-    fn map(&self) -> &Map {
+    fn map(&self) -> &Map<'tcx> {
         &self.map
     }
 
@@ -335,7 +335,7 @@ impl<'tcx> ValueAnalysis<'tcx> for ConstAnalysis<'_, 'tcx> {
 }
 
 impl<'a, 'tcx> ConstAnalysis<'a, 'tcx> {
-    pub fn new(tcx: TyCtxt<'tcx>, body: &'a Body<'tcx>, map: Map) -> Self {
+    pub fn new(tcx: TyCtxt<'tcx>, body: &'a Body<'tcx>, map: Map<'tcx>) -> Self {
         let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id());
         Self {
             map,
@@ -565,12 +565,13 @@ impl<'tcx, 'locals> Collector<'tcx, 'locals> {
         Self { patch: Patch::new(tcx), local_decls }
     }
 
+    #[instrument(level = "trace", skip(self, ecx, map), ret)]
     fn try_make_constant(
         &self,
         ecx: &mut InterpCx<'tcx, DummyMachine>,
         place: Place<'tcx>,
         state: &State<FlatSet<Scalar>>,
-        map: &Map,
+        map: &Map<'tcx>,
     ) -> Option<Const<'tcx>> {
         let ty = place.ty(self.local_decls, self.patch.tcx).ty;
         let layout = ecx.layout_of(ty).ok()?;
@@ -603,10 +604,11 @@ impl<'tcx, 'locals> Collector<'tcx, 'locals> {
     }
 }
 
+#[instrument(level = "trace", skip(map), ret)]
 fn propagatable_scalar(
     place: PlaceIndex,
     state: &State<FlatSet<Scalar>>,
-    map: &Map,
+    map: &Map<'_>,
 ) -> Option<Scalar> {
     if let FlatSet::Elem(value) = state.get_idx(place, map)
         && value.try_to_scalar_int().is_ok()
@@ -618,14 +620,14 @@ fn propagatable_scalar(
     }
 }
 
-#[instrument(level = "trace", skip(ecx, state, map))]
+#[instrument(level = "trace", skip(ecx, state, map), ret)]
 fn try_write_constant<'tcx>(
     ecx: &mut InterpCx<'tcx, DummyMachine>,
     dest: &PlaceTy<'tcx>,
     place: PlaceIndex,
     ty: Ty<'tcx>,
     state: &State<FlatSet<Scalar>>,
-    map: &Map,
+    map: &Map<'tcx>,
 ) -> InterpResult<'tcx> {
     let layout = ecx.layout_of(ty)?;
 
@@ -724,6 +726,7 @@ impl<'mir, 'tcx>
 {
     type FlowState = State<FlatSet<Scalar>>;
 
+    #[instrument(level = "trace", skip(self, results, statement))]
     fn visit_statement_before_primary_effect(
         &mut self,
         results: &mut Results<'tcx, ValueAnalysisWrapper<ConstAnalysis<'_, 'tcx>>>,
@@ -745,6 +748,7 @@ impl<'mir, 'tcx>
         }
     }
 
+    #[instrument(level = "trace", skip(self, results, statement))]
     fn visit_statement_after_primary_effect(
         &mut self,
         results: &mut Results<'tcx, ValueAnalysisWrapper<ConstAnalysis<'_, 'tcx>>>,
@@ -839,7 +843,7 @@ struct OperandCollector<'tcx, 'map, 'locals, 'a> {
     state: &'a State<FlatSet<Scalar>>,
     visitor: &'a mut Collector<'tcx, 'locals>,
     ecx: &'map mut InterpCx<'tcx, DummyMachine>,
-    map: &'map Map,
+    map: &'map Map<'tcx>,
 }
 
 impl<'tcx> Visitor<'tcx> for OperandCollector<'tcx, '_, '_, '_> {

compiler/rustc_mir_transform/src/dataflow_const_prop.rs

@@ -157,7 +157,7 @@ struct TOFinder<'tcx, 'a> {
     param_env: ty::ParamEnv<'tcx>,
     ecx: InterpCx<'tcx, DummyMachine>,
     body: &'a Body<'tcx>,
-    map: &'a Map,
+    map: &'a Map<'tcx>,
     loop_headers: &'a BitSet<BasicBlock>,
     /// We use an arena to avoid cloning the slices when cloning `state`.
     arena: &'a DroplessArena,

compiler/rustc_mir_transform/src/jump_threading.rs

@@ -0,0 +1,65 @@
+- // MIR for `foo` before DataflowConstProp
++ // MIR for `foo` after DataflowConstProp
+  
+  fn foo() -> u32 {
+      let mut _0: u32;
+      let _1: (u32, u32);
+      let mut _4: bool;
+      let mut _5: u32;
+      scope 1 {
+          debug a => _1;
+          let _2: (u32, u32);
+          scope 2 {
+              debug b => _2;
+              let _3: u32;
+              scope 3 {
+                  debug c => _3;
+              }
+          }
+      }
+  
+      bb0: {
+          StorageLive(_1);
+          _1 = const Foo;
+          StorageLive(_2);
+-         _2 = _1;
++         _2 = const (5_u32, 3_u32);
+          StorageLive(_3);
+-         _3 = (_2.1: u32);
++         _3 = const 3_u32;
+          StorageLive(_4);
+          StorageLive(_5);
+-         _5 = _3;
+-         _4 = Ge(move _5, const 2_u32);
+-         switchInt(move _4) -> [0: bb2, otherwise: bb1];
++         _5 = const 3_u32;
++         _4 = const true;
++         switchInt(const true) -> [0: bb2, otherwise: bb1];
+      }
+  
+      bb1: {
+          StorageDead(_5);
+-         _0 = (_2.0: u32);
++         _0 = const 5_u32;
+          goto -> bb3;
+      }
+  
+      bb2: {
+          StorageDead(_5);
+          _0 = const 13_u32;
+          goto -> bb3;
+      }
+  
+      bb3: {
+          StorageDead(_4);
+          StorageDead(_3);
+          StorageDead(_2);
+          StorageDead(_1);
+          return;
+      }
++ }
++ 
++ ALLOC0 (size: 8, align: 4) {
++     05 00 00 00 03 00 00 00                         │ ........
+  }
+  

tests/mir-opt/dataflow-const-prop/aggregate_copy.foo.DataflowConstProp.diff

@@ -0,0 +1,42 @@
+//! Verify that we manage to propagate the value of aggregate `a` even without directly mentioning
+//! the contained scalars.
+//@ test-mir-pass: DataflowConstProp
+
+const Foo: (u32, u32) = (5, 3);
+
+fn foo() -> u32 {
+    // CHECK-LABEL: fn foo(
+    // CHECK: debug a => [[a:_.*]];
+    // CHECK: debug b => [[b:_.*]];
+    // CHECK: debug c => [[c:_.*]];
+
+    // CHECK:bb0: {
+    // CHECK:    [[a]] = const Foo;
+    // CHECK:    [[b]] = const (5_u32, 3_u32);
+    // CHECK:    [[c]] = const 3_u32;
+    // CHECK:    {{_.*}} = const 3_u32;
+    // CHECK:    {{_.*}} = const true;
+    // CHECK:    switchInt(const true) -> [0: bb2, otherwise: bb1];
+
+    // CHECK:bb1: {
+    // CHECK:    _0 = const 5_u32;
+    // CHECK:    goto -> bb3;
+
+    // CHECK:bb2: {
+    // CHECK:    _0 = const 13_u32;
+    // CHECK:    goto -> bb3;
+
+    let a = Foo;
+    // This copies the struct in `a`. We want to ensure that we do track the contents of `a`
+    // because we need to read `b` later.
+    let b = a;
+    let c = b.1;
+    if c >= 2 { b.0 } else { 13 }
+}
+
+fn main() {
+    // CHECK-LABEL: fn main(
+    foo();
+}
+
+// EMIT_MIR aggregate_copy.foo.DataflowConstProp.diff

tests/mir-opt/dataflow-const-prop/aggregate_copy.rs

@@ -106,8 +106,7 @@
 -         _7 = (_10.0: f32);
 +         _7 = const 4f32;
           StorageLive(_8);
--         _8 = (_10.1: std::option::Option<S>);
-+         _8 = const Option::<S>::Some(S(1_i32));
+          _8 = (_10.1: std::option::Option<S>);
           StorageLive(_9);
           _9 = (_10.2: &[f32]);
           StorageDead(_10);
@@ -157,8 +156,7 @@
 +         _23 = const 82f32;
           StorageLive(_24);
           _37 = deref_copy (*_26);
--         _24 = ((*_37).1: std::option::Option<S>);
-+         _24 = const Option::<S>::Some(S(35_i32));
+          _24 = ((*_37).1: std::option::Option<S>);
           StorageLive(_25);
           _38 = deref_copy (*_26);
           _25 = ((*_38).2: &[f32]);
@@ -168,12 +166,11 @@
 -         _28 = _23;
 +         _28 = const 82f32;
           StorageLive(_29);
--         _29 = _24;
-+         _29 = const Option::<S>::Some(S(35_i32));
+          _29 = _24;
           StorageLive(_30);
           _30 = _25;
 -         _27 = BigStruct(move _28, move _29, move _30);
-+         _27 = BigStruct(const 82f32, const Option::<S>::Some(S(35_i32)), move _30);
++         _27 = BigStruct(const 82f32, move _29, move _30);
           StorageDead(_30);
           StorageDead(_29);
           StorageDead(_28);
@@ -199,23 +196,7 @@
       }
 + }
 + 
-+ ALLOC6 (size: 8, align: 4) {
-+     01 00 00 00 23 00 00 00                         │ ....#...
-+ }
-+ 
-+ ALLOC7 (size: 8, align: 4) {
-+     01 00 00 00 23 00 00 00                         │ ....#...
-+ }
-+ 
-+ ALLOC8 (size: 8, align: 4) {
-+     01 00 00 00 23 00 00 00                         │ ....#...
-+ }
-+ 
-+ ALLOC9 (size: 8, align: 4) {
-+     01 00 00 00 01 00 00 00                         │ ........
-+ }
-+ 
-+ ALLOC10 (size: 4, align: 4) {
++ ALLOC6 (size: 4, align: 4) {
 +     01 00 00 00                                     │ ....
   }
 

tests/mir-opt/dataflow-const-prop/struct.main.DataflowConstProp.32bit.diff

@@ -106,8 +106,7 @@
 -         _7 = (_10.0: f32);
 +         _7 = const 4f32;
           StorageLive(_8);
--         _8 = (_10.1: std::option::Option<S>);
-+         _8 = const Option::<S>::Some(S(1_i32));
+          _8 = (_10.1: std::option::Option<S>);
           StorageLive(_9);
           _9 = (_10.2: &[f32]);
           StorageDead(_10);
@@ -157,8 +156,7 @@
 +         _23 = const 82f32;
           StorageLive(_24);
           _37 = deref_copy (*_26);
--         _24 = ((*_37).1: std::option::Option<S>);
-+         _24 = const Option::<S>::Some(S(35_i32));
+          _24 = ((*_37).1: std::option::Option<S>);
           StorageLive(_25);
           _38 = deref_copy (*_26);
           _25 = ((*_38).2: &[f32]);
@@ -168,12 +166,11 @@
 -         _28 = _23;
 +         _28 = const 82f32;
           StorageLive(_29);
--         _29 = _24;
-+         _29 = const Option::<S>::Some(S(35_i32));
+          _29 = _24;
           StorageLive(_30);
           _30 = _25;
 -         _27 = BigStruct(move _28, move _29, move _30);
-+         _27 = BigStruct(const 82f32, const Option::<S>::Some(S(35_i32)), move _30);
++         _27 = BigStruct(const 82f32, move _29, move _30);
           StorageDead(_30);
           StorageDead(_29);
           StorageDead(_28);
@@ -199,23 +196,7 @@
       }
 + }
 + 
-+ ALLOC6 (size: 8, align: 4) {
-+     01 00 00 00 23 00 00 00                         │ ....#...
-+ }
-+ 
-+ ALLOC7 (size: 8, align: 4) {
-+     01 00 00 00 23 00 00 00                         │ ....#...
-+ }
-+ 
-+ ALLOC8 (size: 8, align: 4) {
-+     01 00 00 00 23 00 00 00                         │ ....#...
-+ }
-+ 
-+ ALLOC9 (size: 8, align: 4) {
-+     01 00 00 00 01 00 00 00                         │ ........
-+ }
-+ 
-+ ALLOC10 (size: 4, align: 4) {
++ ALLOC6 (size: 4, align: 4) {
 +     01 00 00 00                                     │ ....
   }
 

tests/mir-opt/dataflow-const-prop/struct.main.DataflowConstProp.64bit.diff

@@ -45,7 +45,7 @@ fn main() {
     const SMALL_VAL: SmallStruct = SmallStruct(4., Some(S(1)), &[]);
 
     // CHECK: [[a1]] = const 4f32;
-    // CHECK: [[b1]] = const Option::<S>::Some(S(1_i32));
+    // CHECK: [[b1]] = ({{_.*}}.1: std::option::Option<S>);
     // CHECK: [[c1]] = ({{_.*}}.2: &[f32]);
     let SmallStruct(a1, b1, c1) = SMALL_VAL;
 
@@ -68,12 +68,12 @@ fn main() {
 
     static BIG_STAT: &BigStruct = &BigStruct(82., Some(S(35)), &[45., 72.]);
     // CHECK: [[a4]] = const 82f32;
-    // CHECK: [[b4]] = const Option::<S>::Some(S(35_i32));
+    // CHECK: [[b4]] = ((*{{_.*}}).1: std::option::Option<S>);
     // CHECK: [[c4]] = ((*{{_.*}}).2: &[f32]);
     let BigStruct(a4, b4, c4) = *BIG_STAT;
 
     // We arbitrarily limit the size of synthetized values to 4 pointers.
     // `BigStruct` can be read, but we will keep a MIR aggregate for this.
-    // CHECK: [[bs]] = BigStruct(const 82f32, const Option::<S>::Some(S(35_i32)), move {{_.*}});
+    // CHECK: [[bs]] = BigStruct(const 82f32, move {{.*}}, move {{_.*}});
     let bs = BigStruct(a4, b4, c4);
 }

tests/mir-opt/dataflow-const-prop/struct.rs

@@ -7,10 +7,20 @@
 
       bb0: {
           StorageLive(_1);
-          _1 = Less;
-          _0 = move _1 as i8 (Transmute);
+-         _1 = Less;
+-         _0 = move _1 as i8 (Transmute);
++         _1 = const Less;
++         _0 = const std::cmp::Ordering::Less as i8 (Transmute);
           StorageDead(_1);
           return;
       }
++ }
++ 
++ ALLOC0 (size: 1, align: 1) {
++     ff                                              │ .
++ }
++ 
++ ALLOC1 (size: 1, align: 1) {
++     ff                                              │ .
   }