Propagate places through assignments.

cjgillot · cjgillot · commit 82b6247b8c94 · 2024-07-04T06:29:20.000Z
diff --git a/compiler/rustc_mir_dataflow/src/value_analysis.rs b/compiler/rustc_mir_dataflow/src/value_analysis.rs
@@ -36,7 +36,7 @@ use std::fmt::{Debug, Formatter};
 use std::ops::Range;
 
 use rustc_data_structures::captures::Captures;
-use rustc_data_structures::fx::{FxHashMap, StdEntry};
+use rustc_data_structures::fx::{FxHashMap, FxIndexSet, StdEntry};
 use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_index::bit_set::BitSet;
 use rustc_index::IndexVec;
@@ -796,7 +796,52 @@ impl<'tcx> Map<'tcx> {
             self.locals[local] = Some(place);
         }
 
-        PlaceCollector { tcx, body, map: self }.visit_body(body);
+        // Collect syntactic places and assignments between them.
+        let mut collector =
+            PlaceCollector { tcx, body, map: self, assignments: Default::default() };
+        collector.visit_body(body);
+        let PlaceCollector { mut assignments, .. } = collector;
+
+        // Just collecting syntactic places is not enough. We may need to propagate this pattern:
+        //      _1 = (const 5u32, const 13i64);
+        //      _2 = _1;
+        //      _3 = (_2.0 as u32);
+        //
+        // `_1.0` does not appear, but we still need to track it. This is achieved by propagating
+        // projections from assignments. We recorded an assignment between `_2` and `_1`, so we
+        // want `_1` and `_2` to have the same sub-places.
+        //
+        // This is what this fixpoint loop does. While we are still creating places, run through
+        // all the assignments, and register places for children.
+        let mut num_places = 0;
+        while num_places < self.places.len() {
+            num_places = self.places.len();
+
+            for assign in 0.. {
+                let Some(&(lhs, rhs)) = assignments.get_index(assign) else { break };
+
+                // Mirror children from `lhs` in `rhs`.
+                let mut child = self.places[lhs].first_child;
+                while let Some(lhs_child) = child {
+                    let PlaceInfo { ty, proj_elem, next_sibling, .. } = self.places[lhs_child];
+                    let rhs_child =
+                        self.register_place(ty, rhs, proj_elem.expect("child is not a projection"));
+                    assignments.insert((lhs_child, rhs_child));
+                    child = next_sibling;
+                }
+
+                // Conversely, mirror children from `rhs` in `lhs`.
+                let mut child = self.places[rhs].first_child;
+                while let Some(rhs_child) = child {
+                    let PlaceInfo { ty, proj_elem, next_sibling, .. } = self.places[rhs_child];
+                    let lhs_child =
+                        self.register_place(ty, lhs, proj_elem.expect("child is not a projection"));
+                    assignments.insert((lhs_child, rhs_child));
+                    child = next_sibling;
+                }
+            }
+        }
+        drop(assignments);
 
         // Create values for places whose type have scalar layout.
         let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id());
@@ -879,17 +924,14 @@ struct PlaceCollector<'a, 'b, 'tcx> {
     tcx: TyCtxt<'tcx>,
     body: &'b Body<'tcx>,
     map: &'a mut Map<'tcx>,
+    assignments: FxIndexSet<(PlaceIndex, PlaceIndex)>,
 }
 
-impl<'tcx> Visitor<'tcx> for PlaceCollector<'_, '_, 'tcx> {
+impl<'tcx> PlaceCollector<'_, '_, 'tcx> {
     #[tracing::instrument(level = "trace", skip(self))]
-    fn visit_place(&mut self, place: &Place<'tcx>, ctxt: PlaceContext, _: Location) {
-        if !ctxt.is_use() {
-            return;
-        }
-
+    fn register_place(&mut self, place: Place<'tcx>) -> Option<PlaceIndex> {
         // Create a place for this projection.
-        let Some(mut place_index) = self.map.locals[place.local] else { return };
+        let mut place_index = self.map.locals[place.local]?;
         let mut ty = PlaceTy::from_ty(self.body.local_decls[place.local].ty);
         tracing::trace!(?place_index, ?ty);
 
@@ -903,7 +945,7 @@ impl<'tcx> Visitor<'tcx> for PlaceCollector<'_, '_, 'tcx> {
         }
 
         for proj in place.projection {
-            let Ok(track_elem) = proj.try_into() else { return };
+            let track_elem = proj.try_into().ok()?;
             ty = ty.projection_ty(self.tcx, proj);
             place_index = self.map.register_place(ty.ty, place_index, track_elem);
             tracing::trace!(?proj, ?place_index, ?ty);
@@ -917,6 +959,63 @@ impl<'tcx> Visitor<'tcx> for PlaceCollector<'_, '_, 'tcx> {
                 self.map.register_place(discriminant_ty, place_index, TrackElem::Discriminant);
             }
         }
+
+        Some(place_index)
+    }
+}
+
+impl<'tcx> Visitor<'tcx> for PlaceCollector<'_, '_, 'tcx> {
+    #[tracing::instrument(level = "trace", skip(self))]
+    fn visit_place(&mut self, place: &Place<'tcx>, ctxt: PlaceContext, _: Location) {
+        if !ctxt.is_use() {
+            return;
+        }
+
+        self.register_place(*place);
+    }
+
+    fn visit_assign(&mut self, lhs: &Place<'tcx>, rhs: &Rvalue<'tcx>, location: Location) {
+        self.super_assign(lhs, rhs, location);
+
+        match rhs {
+            Rvalue::Use(Operand::Move(rhs) | Operand::Copy(rhs)) | Rvalue::CopyForDeref(rhs) => {
+                let Some(lhs) = self.register_place(*lhs) else { return };
+                let Some(rhs) = self.register_place(*rhs) else { return };
+                self.assignments.insert((lhs, rhs));
+            }
+            Rvalue::Aggregate(kind, fields) => {
+                let Some(mut lhs) = self.register_place(*lhs) else { return };
+                match **kind {
+                    // Do not propagate unions.
+                    AggregateKind::Adt(_, _, _, _, Some(_)) => return,
+                    AggregateKind::Adt(_, variant, _, _, None) => {
+                        let ty = self.map.places[lhs].ty;
+                        if ty.is_enum() {
+                            lhs = self.map.register_place(ty, lhs, TrackElem::Variant(variant));
+                        }
+                    }
+                    AggregateKind::RawPtr(..)
+                    | AggregateKind::Array(_)
+                    | AggregateKind::Tuple
+                    | AggregateKind::Closure(..)
+                    | AggregateKind::Coroutine(..)
+                    | AggregateKind::CoroutineClosure(..) => {}
+                }
+                for (index, field) in fields.iter_enumerated() {
+                    if let Some(rhs) = field.place()
+                        && let Some(rhs) = self.register_place(rhs)
+                    {
+                        let lhs = self.map.register_place(
+                            self.map.places[rhs].ty,
+                            lhs,
+                            TrackElem::Field(index),
+                        );
+                        self.assignments.insert((lhs, rhs));
+                    }
+                }
+            }
+            _ => {}
+        }
     }
 }
 
diff --git a/tests/mir-opt/dataflow-const-prop/aggregate_copy.foo.DataflowConstProp.diff b/tests/mir-opt/dataflow-const-prop/aggregate_copy.foo.DataflowConstProp.diff
@@ -22,19 +22,25 @@
           StorageLive(_1);
           _1 = const Foo;
           StorageLive(_2);
-          _2 = _1;
+-         _2 = _1;
++         _2 = const (5_u32, 3_u32);
           StorageLive(_3);
-          _3 = (_2.1: u32);
+-         _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 = _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 = (_2.0: u32);
++         _0 = const 5_u32;
           goto -> bb3;
       }
   
@@ -51,5 +57,9 @@
           StorageDead(_1);
           return;
       }
++ }
++ 
++ ALLOC0 (size: 8, align: 4) {
++     05 00 00 00 03 00 00 00                         │ ........
   }
   
diff --git a/tests/mir-opt/dataflow-const-prop/aggregate_copy.rs b/tests/mir-opt/dataflow-const-prop/aggregate_copy.rs
@@ -12,12 +12,14 @@ fn foo() -> u32 {
 
     // CHECK:bb0: {
     // CHECK:    [[a]] = const Foo;
-    // CHECK:    [[b]] = [[a]];
-    // CHECK:    [[c]] = ([[b]].1: u32);
-    // CHECK:    switchInt(move {{_.*}}) -> [0: bb2, otherwise: bb1];
+    // 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 = ([[b]].0: u32);
+    // CHECK:    _0 = const 5_u32;
     // CHECK:    goto -> bb3;
 
     // CHECK:bb2: {
diff --git a/tests/mir-opt/dataflow-const-prop/repr_transparent.main.DataflowConstProp.diff b/tests/mir-opt/dataflow-const-prop/repr_transparent.main.DataflowConstProp.diff
@@ -32,7 +32,7 @@
           StorageDead(_5);
           StorageDead(_4);
 -         _2 = I32(move _3);
-+         _2 = I32(const 0_i32);
++         _2 = const I32(0_i32);
           StorageDead(_3);
           _0 = const ();
           StorageDead(_2);
@@ -42,6 +42,10 @@
 + }
 + 
 + ALLOC0 (size: 4, align: 4) {
++     00 00 00 00                                     │ ....
++ }
++ 
++ ALLOC1 (size: 4, align: 4) {
 +     00 00 00 00                                     │ ....
   }
   
diff --git a/tests/mir-opt/dataflow-const-prop/repr_transparent.rs b/tests/mir-opt/dataflow-const-prop/repr_transparent.rs
@@ -15,6 +15,6 @@ fn main() {
     // CHECK: [[x]] = const I32(0_i32);
     let x = I32(0);
 
-    // CHECK: [[y]] = I32(const 0_i32);
+    // CHECK: [[y]] = const I32(0_i32);
     let y = I32(x.0 + x.0);
 }
diff --git a/tests/mir-opt/dataflow-const-prop/transmute.less_as_i8.DataflowConstProp.32bit.diff b/tests/mir-opt/dataflow-const-prop/transmute.less_as_i8.DataflowConstProp.32bit.diff
@@ -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                                              │ .
   }
   
diff --git a/tests/mir-opt/dataflow-const-prop/transmute.less_as_i8.DataflowConstProp.64bit.diff b/tests/mir-opt/dataflow-const-prop/transmute.less_as_i8.DataflowConstProp.64bit.diff
@@ -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                                              │ .
   }
   
diff --git a/tests/mir-opt/dataflow-const-prop/tuple.main.DataflowConstProp.32bit.diff b/tests/mir-opt/dataflow-const-prop/tuple.main.DataflowConstProp.32bit.diff
@@ -81,7 +81,7 @@
 -         _14 = _6;
 -         _11 = (move _12, move _13, move _14);
 +         _14 = const 11_i32;
-+         _11 = (const 6_i32, move _13, const 11_i32);
++         _11 = (const 6_i32, const (2_i32, 3_i32), const 11_i32);
           StorageDead(_14);
           StorageDead(_13);
           StorageDead(_12);
@@ -103,6 +103,10 @@
 + }
 + 
 + ALLOC2 (size: 8, align: 4) {
++     02 00 00 00 03 00 00 00                         │ ........
++ }
++ 
++ ALLOC3 (size: 8, align: 4) {
 +     01 00 00 00 02 00 00 00                         │ ........
   }
   
diff --git a/tests/mir-opt/dataflow-const-prop/tuple.main.DataflowConstProp.64bit.diff b/tests/mir-opt/dataflow-const-prop/tuple.main.DataflowConstProp.64bit.diff
@@ -81,7 +81,7 @@
 -         _14 = _6;
 -         _11 = (move _12, move _13, move _14);
 +         _14 = const 11_i32;
-+         _11 = (const 6_i32, move _13, const 11_i32);
++         _11 = (const 6_i32, const (2_i32, 3_i32), const 11_i32);
           StorageDead(_14);
           StorageDead(_13);
           StorageDead(_12);
@@ -103,6 +103,10 @@
 + }
 + 
 + ALLOC2 (size: 8, align: 4) {
++     02 00 00 00 03 00 00 00                         │ ........
++ }
++ 
++ ALLOC3 (size: 8, align: 4) {
 +     01 00 00 00 02 00 00 00                         │ ........
   }
   
diff --git a/tests/mir-opt/dataflow-const-prop/tuple.rs b/tests/mir-opt/dataflow-const-prop/tuple.rs
@@ -22,7 +22,6 @@ fn main() {
     // CHECK: [[c]] = const 11_i32;
     let c = a.0 + a.1 + b;
 
-    // CHECK: [[a2:_.*]] = const (2_i32, 3_i32);
-    // CHECK: [[d]] = (const 6_i32, move [[a2]], const 11_i32);
+    // CHECK: [[d]] = (const 6_i32, const (2_i32, 3_i32), const 11_i32);
     let d = (b, a, c);
 }
diff --git a/tests/mir-opt/jump_threading.aggregate_copy.JumpThreading.panic-abort.diff b/tests/mir-opt/jump_threading.aggregate_copy.JumpThreading.panic-abort.diff
@@ -29,7 +29,8 @@
           StorageLive(_5);
           _5 = _3;
           _4 = Eq(move _5, const 2_u32);
-          switchInt(move _4) -> [0: bb2, otherwise: bb1];
+-         switchInt(move _4) -> [0: bb2, otherwise: bb1];
++         goto -> bb2;
       }
   
       bb1: {
diff --git a/tests/mir-opt/jump_threading.aggregate_copy.JumpThreading.panic-unwind.diff b/tests/mir-opt/jump_threading.aggregate_copy.JumpThreading.panic-unwind.diff
@@ -29,7 +29,8 @@
           StorageLive(_5);
           _5 = _3;
           _4 = Eq(move _5, const 2_u32);
-          switchInt(move _4) -> [0: bb2, otherwise: bb1];
+-         switchInt(move _4) -> [0: bb2, otherwise: bb1];
++         goto -> bb2;
       }
   
       bb1: {
diff --git a/tests/mir-opt/jump_threading.rs b/tests/mir-opt/jump_threading.rs
@@ -509,7 +509,7 @@ fn assume(a: u8, b: bool) -> u8 {
 /// Verify that jump threading succeeds seeing through copies of aggregates.
 fn aggregate_copy() -> u32 {
     // CHECK-LABEL: fn aggregate_copy(
-    // CHECK: switchInt(
+    // CHECK-NOT: switchInt(
 
     const Foo: (u32, u32) = (5, 3);
 

Original file line number	Diff line number	Diff line change
`@@ -15,6 +15,6 @@ fn main() {`
`15`	`15`	`// CHECK: [[x]] = const I32(0_i32);`
`16`	`16`	`let x = I32(0);`
`17`	`17`
`18`		`- // CHECK: [[y]] = I32(const 0_i32);`
	`18`	`+ // CHECK: [[y]] = const I32(0_i32);`
`19`	`19`	`let y = I32(x.0 + x.0);`
`20`	`20`	`}`
Original file line number	Diff line number	Diff line change
`@@ -22,7 +22,6 @@ fn main() {`
`22`	`22`	`// CHECK: [[c]] = const 11_i32;`
`23`	`23`	`let c = a.0 + a.1 + b;`
`24`	`24`
`25`		`- // CHECK: [[a2:_.*]] = const (2_i32, 3_i32);`
`26`		`- // CHECK: [[d]] = (const 6_i32, move [[a2]], const 11_i32);`
	`25`	`+ // CHECK: [[d]] = (const 6_i32, const (2_i32, 3_i32), const 11_i32);`
`27`	`26`	`let d = (b, a, c);`
`28`	`27`	`}`
Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,8 @@`
`29`	`29`	`StorageLive(_5);`
`30`	`30`	`_5 = _3;`
`31`	`31`	`_4 = Eq(move _5, const 2_u32);`
`32`		`- switchInt(move _4) -> [0: bb2, otherwise: bb1];`
	`32`	`+- switchInt(move _4) -> [0: bb2, otherwise: bb1];`
	`33`	`++ goto -> bb2;`
`33`	`34`	`}`
`34`	`35`
`35`	`36`	`bb1: {`