Rollup merge of #108208 - cjgillot:flood-enum, r=oli-obk

Correctly handle aggregates in DataflowConstProp

The previous implementation from https://github.com/rust-lang/rust/pull/107411 flooded target of an aggregate assignment with `Bottom`, corresponding to the `deinit` that the interpreter does.

As a consequence, when assigning `target = Enum::Variant#i(...)` all the `(target as Variant#j)` were at `Bottom` while they should have been `Top`.

This PR replaces that flooding with `Top`.

Aside, it corrects a second bug where the wrong place would be used to assign to enum variant fields, resulting to nothing happening.

Fixes https://github.com/rust-lang/rust/issues/108166

4 files changed, 120 insertions, 9 deletions

diff --git a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
index 19019e3ef74..49ded10ba1f 100644
--- a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
+++ b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
@@ -122,7 +122,10 @@ impl<'tcx> ValueAnalysis<'tcx> for ConstAnalysis<'_, 'tcx> {
     ) {
         match rvalue {
             Rvalue::Aggregate(kind, operands) => {
-                state.flood_with(target.as_ref(), self.map(), FlatSet::Bottom);
+                // If we assign `target = Enum::Variant#0(operand)`,
+                // we must make sure that all `target as Variant#i` are `Top`.
+                state.flood(target.as_ref(), self.map());
+
                 if let Some(target_idx) = self.map().find(target.as_ref()) {
                     let (variant_target, variant_index) = match **kind {
                         AggregateKind::Tuple | AggregateKind::Closure(..) => {
@@ -131,18 +134,21 @@ impl<'tcx> ValueAnalysis<'tcx> for ConstAnalysis<'_, 'tcx> {
                         AggregateKind::Adt(def_id, variant_index, ..) => {
                             match self.tcx.def_kind(def_id) {
                                 DefKind::Struct => (Some(target_idx), None),
-                                DefKind::Enum => (Some(target_idx), Some(variant_index)),
+                                DefKind::Enum => (
+                                    self.map.apply(target_idx, TrackElem::Variant(variant_index)),
+                                    Some(variant_index),
+                                ),
                                 _ => (None, None),
                             }
                         }
                         _ => (None, None),
                     };
-                    if let Some(target) = variant_target {
+                    if let Some(variant_target_idx) = variant_target {
                         for (field_index, operand) in operands.iter().enumerate() {
-                            if let Some(field) = self
-                                .map()
-                                .apply(target, TrackElem::Field(Field::from_usize(field_index)))
-                            {
+                            if let Some(field) = self.map().apply(
+                                variant_target_idx,
+                                TrackElem::Field(Field::from_usize(field_index)),
+                            ) {
                                 let result = self.handle_operand(operand, state);
                                 state.insert_idx(field, result, self.map());
                             }
@@ -151,6 +157,11 @@ impl<'tcx> ValueAnalysis<'tcx> for ConstAnalysis<'_, 'tcx> {
                     if let Some(variant_index) = variant_index
                         && let Some(discr_idx) = self.map().apply(target_idx, TrackElem::Discriminant)
                     {
+                        // We are assigning the discriminant as part of an aggregate.
+                        // This discriminant can only alias a variant field's value if the operand
+                        // had an invalid value for that type.
+                        // Using invalid values is UB, so we are allowed to perform the assignment
+                        // without extra flooding.
                         let enum_ty = target.ty(self.local_decls, self.tcx).ty;
                         if let Some(discr_val) = self.eval_discriminant(enum_ty, variant_index) {
                             state.insert_value_idx(discr_idx, FlatSet::Elem(discr_val), &self.map);
diff --git a/tests/mir-opt/dataflow-const-prop/enum.multiple.DataflowConstProp.diff b/tests/mir-opt/dataflow-const-prop/enum.multiple.DataflowConstProp.diff
new file mode 100644
index 00000000000..c4002d65e5d
--- /dev/null
+++ b/tests/mir-opt/dataflow-const-prop/enum.multiple.DataflowConstProp.diff
@@ -0,0 +1,82 @@
+- // MIR for `multiple` before DataflowConstProp
++ // MIR for `multiple` after DataflowConstProp
+  
+  fn multiple(_1: bool, _2: u8) -> () {
+      debug x => _1;                       // in scope 0 at $DIR/enum.rs:+0:13: +0:14
+      debug i => _2;                       // in scope 0 at $DIR/enum.rs:+0:22: +0:23
+      let mut _0: ();                      // return place in scope 0 at $DIR/enum.rs:+0:29: +0:29
+      let _3: std::option::Option<u8>;     // in scope 0 at $DIR/enum.rs:+1:9: +1:10
+      let mut _4: bool;                    // in scope 0 at $DIR/enum.rs:+1:16: +1:17
+      let mut _5: u8;                      // in scope 0 at $DIR/enum.rs:+2:14: +2:15
+      let mut _7: isize;                   // in scope 0 at $DIR/enum.rs:+9:23: +9:30
+      scope 1 {
+          debug e => _3;                   // in scope 1 at $DIR/enum.rs:+1:9: +1:10
+          let _6: u8;                      // in scope 1 at $DIR/enum.rs:+9:9: +9:10
+          let _8: u8;                      // in scope 1 at $DIR/enum.rs:+9:28: +9:29
+          scope 2 {
+              debug x => _6;               // in scope 2 at $DIR/enum.rs:+9:9: +9:10
+              let _9: u8;                  // in scope 2 at $DIR/enum.rs:+11:9: +11:10
+              scope 4 {
+                  debug y => _9;           // in scope 4 at $DIR/enum.rs:+11:9: +11:10
+              }
+          }
+          scope 3 {
+              debug i => _8;               // in scope 3 at $DIR/enum.rs:+9:28: +9:29
+          }
+      }
+  
+      bb0: {
+          StorageLive(_3);                 // scope 0 at $DIR/enum.rs:+1:9: +1:10
+          StorageLive(_4);                 // scope 0 at $DIR/enum.rs:+1:16: +1:17
+          _4 = _1;                         // scope 0 at $DIR/enum.rs:+1:16: +1:17
+          switchInt(move _4) -> [0: bb2, otherwise: bb1]; // scope 0 at $DIR/enum.rs:+1:16: +1:17
+      }
+  
+      bb1: {
+          StorageLive(_5);                 // scope 0 at $DIR/enum.rs:+2:14: +2:15
+          _5 = _2;                         // scope 0 at $DIR/enum.rs:+2:14: +2:15
+          _3 = Option::<u8>::Some(move _5); // scope 0 at $DIR/enum.rs:+2:9: +2:16
+          StorageDead(_5);                 // scope 0 at $DIR/enum.rs:+2:15: +2:16
+          goto -> bb3;                     // scope 0 at $DIR/enum.rs:+1:13: +5:6
+      }
+  
+      bb2: {
+          _3 = Option::<u8>::None;         // scope 0 at $DIR/enum.rs:+4:9: +4:13
+          goto -> bb3;                     // scope 0 at $DIR/enum.rs:+1:13: +5:6
+      }
+  
+      bb3: {
+          StorageDead(_4);                 // scope 0 at $DIR/enum.rs:+5:5: +5:6
+          StorageLive(_6);                 // scope 1 at $DIR/enum.rs:+9:9: +9:10
+          _7 = discriminant(_3);           // scope 1 at $DIR/enum.rs:+9:19: +9:20
+          switchInt(move _7) -> [0: bb4, 1: bb6, otherwise: bb5]; // scope 1 at $DIR/enum.rs:+9:13: +9:20
+      }
+  
+      bb4: {
+          _6 = const 0_u8;                 // scope 1 at $DIR/enum.rs:+9:45: +9:46
+          goto -> bb7;                     // scope 1 at $DIR/enum.rs:+9:45: +9:46
+      }
+  
+      bb5: {
+          unreachable;                     // scope 1 at $DIR/enum.rs:+9:19: +9:20
+      }
+  
+      bb6: {
+          StorageLive(_8);                 // scope 1 at $DIR/enum.rs:+9:28: +9:29
+          _8 = ((_3 as Some).0: u8);       // scope 1 at $DIR/enum.rs:+9:28: +9:29
+          _6 = _8;                         // scope 3 at $DIR/enum.rs:+9:34: +9:35
+          StorageDead(_8);                 // scope 1 at $DIR/enum.rs:+9:34: +9:35
+          goto -> bb7;                     // scope 1 at $DIR/enum.rs:+9:34: +9:35
+      }
+  
+      bb7: {
+          StorageLive(_9);                 // scope 2 at $DIR/enum.rs:+11:9: +11:10
+          _9 = _6;                         // scope 2 at $DIR/enum.rs:+11:13: +11:14
+          _0 = const ();                   // scope 0 at $DIR/enum.rs:+0:29: +12:2
+          StorageDead(_9);                 // scope 2 at $DIR/enum.rs:+12:1: +12:2
+          StorageDead(_6);                 // scope 1 at $DIR/enum.rs:+12:1: +12:2
+          StorageDead(_3);                 // scope 0 at $DIR/enum.rs:+12:1: +12:2
+          return;                          // scope 0 at $DIR/enum.rs:+12:2: +12:2
+      }
+  }
+  
diff --git a/tests/mir-opt/dataflow-const-prop/enum.rs b/tests/mir-opt/dataflow-const-prop/enum.rs
index 7ea405bd9c4..79a20d7ef45 100644
--- a/tests/mir-opt/dataflow-const-prop/enum.rs
+++ b/tests/mir-opt/dataflow-const-prop/enum.rs
@@ -46,7 +46,23 @@ fn mutate_discriminant() -> u8 {
     )
 }
 
+// EMIT_MIR enum.multiple.DataflowConstProp.diff
+fn multiple(x: bool, i: u8) {
+    let e = if x {
+        Some(i)
+    } else {
+        None
+    };
+    // The dataflow state must have:
+    //   discriminant(e) => Top
+    //   (e as Some).0 => Top
+    let x = match e { Some(i) => i, None => 0 };
+    // Therefore, `x` should be `Top` here, and no replacement shall happen.
+    let y = x;
+}
+
 fn main() {
     simple();
     mutate_discriminant();
+    multiple(false, 5);
 }
diff --git a/tests/mir-opt/dataflow-const-prop/enum.simple.DataflowConstProp.diff b/tests/mir-opt/dataflow-const-prop/enum.simple.DataflowConstProp.diff
index 1fb65e65845..22bdc35d694 100644
--- a/tests/mir-opt/dataflow-const-prop/enum.simple.DataflowConstProp.diff
+++ b/tests/mir-opt/dataflow-const-prop/enum.simple.DataflowConstProp.diff
@@ -45,8 +45,10 @@
   
       bb3: {
           StorageLive(_4);                 // scope 1 at $DIR/enum.rs:+2:29: +2:30
-          _4 = ((_1 as V1).0: i32);        // scope 1 at $DIR/enum.rs:+2:29: +2:30
-          _2 = _4;                         // scope 3 at $DIR/enum.rs:+2:35: +2:36
+-         _4 = ((_1 as V1).0: i32);        // scope 1 at $DIR/enum.rs:+2:29: +2:30
+-         _2 = _4;                         // scope 3 at $DIR/enum.rs:+2:35: +2:36
++         _4 = const 0_i32;                // scope 1 at $DIR/enum.rs:+2:29: +2:30
++         _2 = const 0_i32;                // scope 3 at $DIR/enum.rs:+2:35: +2:36
           StorageDead(_4);                 // scope 1 at $DIR/enum.rs:+2:35: +2:36
           goto -> bb4;                     // scope 1 at $DIR/enum.rs:+2:35: +2:36
       }