Auto merge of #95125 - JakobDegen:uninit-variant-rvalue, r=oli-obk

Add new `Deinit` statement

This rvalue replaces `SetDiscriminant` for ADTs. This PR is an alternative to #94590 , which only specifies that the behavior of `SetDiscriminant` is the same as what this rvalue would do. The motivation for this change are discussed in that PR and [on Zulip](https://rust-lang.zulipchat.com/#narrow/stream/189540-t-compiler.2Fwg-mir-opt/topic/SetDiscriminant.20and.20aggregate.20initialization.20.2394590)

r? `@oli-obk`

compiler/rustc_borrowck/src/dataflow.rs

@@ -386,6 +386,7 @@ impl<'tcx> rustc_mir_dataflow::GenKillAnalysis<'tcx> for Borrows<'_, 'tcx> {
 
             mir::StatementKind::FakeRead(..)
             | mir::StatementKind::SetDiscriminant { .. }
+            | mir::StatementKind::Deinit(..)
             | mir::StatementKind::StorageLive(..)
             | mir::StatementKind::Retag { .. }
             | mir::StatementKind::AscribeUserType(..)

compiler/rustc_borrowck/src/def_use.rs

@@ -72,5 +72,9 @@ pub fn categorize(context: PlaceContext) -> Option<DefUse> {
 
         // Debug info is neither def nor use.
         PlaceContext::NonUse(NonUseContext::VarDebugInfo) => None,
+
+        PlaceContext::MutatingUse(MutatingUseContext::Deinit | MutatingUseContext::SetDiscriminant) => {
+            bug!("These statements are not allowed in this MIR phase")
+        }
     }
 }

compiler/rustc_borrowck/src/invalidation.rs

@@ -63,9 +63,6 @@ impl<'cx, 'tcx> Visitor<'tcx> for InvalidationGenerator<'cx, 'tcx> {
             StatementKind::FakeRead(box (_, _)) => {
                 // Only relevant for initialized/liveness/safety checks.
             }
-            StatementKind::SetDiscriminant { place, variant_index: _ } => {
-                self.mutate_place(location, **place, Shallow(None));
-            }
             StatementKind::CopyNonOverlapping(box rustc_middle::mir::CopyNonOverlapping {
                 ref src,
                 ref dst,
@@ -91,6 +88,9 @@ impl<'cx, 'tcx> Visitor<'tcx> for InvalidationGenerator<'cx, 'tcx> {
                     LocalMutationIsAllowed::Yes,
                 );
             }
+            StatementKind::Deinit(..) | StatementKind::SetDiscriminant { .. } => {
+                bug!("Statement not allowed in this MIR phase")
+            }
         }
 
         self.super_statement(statement, location);

compiler/rustc_borrowck/src/lib.rs

@@ -626,9 +626,6 @@ impl<'cx, 'tcx> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx> for MirBorrowckCtx
                     flow_state,
                 );
             }
-            StatementKind::SetDiscriminant { place, variant_index: _ } => {
-                self.mutate_place(location, (**place, span), Shallow(None), flow_state);
-            }
             StatementKind::CopyNonOverlapping(box rustc_middle::mir::CopyNonOverlapping {
                 ..
             }) => {
@@ -654,6 +651,9 @@ impl<'cx, 'tcx> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx> for MirBorrowckCtx
                     flow_state,
                 );
             }
+            StatementKind::Deinit(..) | StatementKind::SetDiscriminant { .. } => {
+                bug!("Statement not allowed in this MIR phase")
+            }
         }
     }
 

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -1303,28 +1303,6 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                     );
                 }
             }
-            StatementKind::SetDiscriminant { ref place, variant_index } => {
-                let place_type = place.ty(body, tcx).ty;
-                let adt = match place_type.kind() {
-                    ty::Adt(adt, _) if adt.is_enum() => adt,
-                    _ => {
-                        span_bug!(
-                            stmt.source_info.span,
-                            "bad set discriminant ({:?} = {:?}): lhs is not an enum",
-                            place,
-                            variant_index
-                        );
-                    }
-                };
-                if variant_index.as_usize() >= adt.variants().len() {
-                    span_bug!(
-                        stmt.source_info.span,
-                        "bad set discriminant ({:?} = {:?}): value of of range",
-                        place,
-                        variant_index
-                    );
-                };
-            }
             StatementKind::AscribeUserType(box (ref place, ref projection), variance) => {
                 let place_ty = place.ty(body, tcx).ty;
                 if let Err(terr) = self.relate_type_and_user_type(
@@ -1358,6 +1336,9 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             | StatementKind::Retag { .. }
             | StatementKind::Coverage(..)
             | StatementKind::Nop => {}
+            StatementKind::Deinit(..) | StatementKind::SetDiscriminant { .. } => {
+                bug!("Statement not allowed in this MIR phase")
+            }
         }
     }
 

compiler/rustc_codegen_cranelift/src/base.rs

@@ -772,6 +772,7 @@ fn codegen_stmt<'tcx>(
         }
         StatementKind::StorageLive(_)
         | StatementKind::StorageDead(_)
+        | StatementKind::Deinit(_)
         | StatementKind::Nop
         | StatementKind::FakeRead(..)
         | StatementKind::Retag { .. }

compiler/rustc_codegen_cranelift/src/constant.rs

@@ -518,6 +518,7 @@ pub(crate) fn mir_operand_get_const_val<'tcx>(
                         StatementKind::Assign(_)
                         | StatementKind::FakeRead(_)
                         | StatementKind::SetDiscriminant { .. }
+                        | StatementKind::Deinit(_)
                         | StatementKind::StorageLive(_)
                         | StatementKind::StorageDead(_)
                         | StatementKind::Retag(_, _)

compiler/rustc_codegen_ssa/src/mir/analyze.rs

@@ -211,6 +211,8 @@ impl<'mir, 'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> Visitor<'tcx>
 
             PlaceContext::MutatingUse(
                 MutatingUseContext::Store
+                | MutatingUseContext::Deinit
+                | MutatingUseContext::SetDiscriminant
                 | MutatingUseContext::AsmOutput
                 | MutatingUseContext::Borrow
                 | MutatingUseContext::AddressOf

compiler/rustc_codegen_ssa/src/mir/statement.rs

@@ -48,6 +48,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     .codegen_set_discr(&mut bx, variant_index);
                 bx
             }
+            mir::StatementKind::Deinit(..) => {
+                // For now, don't codegen this to anything. In the future it may be worth
+                // experimenting with what kind of information we can emit to LLVM without hurting
+                // perf here
+                bx
+            }
             mir::StatementKind::StorageLive(local) => {
                 if let LocalRef::Place(cg_place) = self.locals[local] {
                     cg_place.storage_live(&mut bx);

compiler/rustc_const_eval/src/interpret/memory.rs

@@ -890,6 +890,11 @@ impl<'tcx, 'a, Tag: Provenance, Extra> AllocRefMut<'a, 'tcx, Tag, Extra> {
     ) -> InterpResult<'tcx> {
         self.write_scalar(alloc_range(offset, self.tcx.data_layout().pointer_size), val)
     }
+
+    /// Mark the entire referenced range as uninitalized
+    pub fn write_uninit(&mut self) {
+        self.alloc.mark_init(self.range, false);
+    }
 }
 
 impl<'tcx, 'a, Tag: Provenance, Extra> AllocRef<'a, 'tcx, Tag, Extra> {

compiler/rustc_const_eval/src/interpret/place.rs

@@ -791,6 +791,42 @@ where
         }
     }
 
+    pub fn write_uninit(&mut self, dest: &PlaceTy<'tcx, M::PointerTag>) -> InterpResult<'tcx> {
+        let mplace = match dest.place {
+            Place::Ptr(mplace) => MPlaceTy { mplace, layout: dest.layout },
+            Place::Local { frame, local } => {
+                match M::access_local_mut(self, frame, local)? {
+                    Ok(local) => match dest.layout.abi {
+                        Abi::Scalar(_) => {
+                            *local = LocalValue::Live(Operand::Immediate(Immediate::Scalar(
+                                ScalarMaybeUninit::Uninit,
+                            )));
+                            return Ok(());
+                        }
+                        Abi::ScalarPair(..) => {
+                            *local = LocalValue::Live(Operand::Immediate(Immediate::ScalarPair(
+                                ScalarMaybeUninit::Uninit,
+                                ScalarMaybeUninit::Uninit,
+                            )));
+                            return Ok(());
+                        }
+                        _ => self.force_allocation(dest)?,
+                    },
+                    Err(mplace) => {
+                        // The local is in memory, go on below.
+                        MPlaceTy { mplace, layout: dest.layout }
+                    }
+                }
+            }
+        };
+        let Some(mut alloc) = self.get_place_alloc_mut(&mplace)? else {
+            // Zero-sized access
+            return Ok(());
+        };
+        alloc.write_uninit();
+        Ok(())
+    }
+
     /// Copies the data from an operand to a place. This does not support transmuting!
     /// Use `copy_op_transmute` if the layouts could disagree.
     #[inline(always)]

compiler/rustc_const_eval/src/interpret/step.rs

@@ -90,6 +90,11 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 self.write_discriminant(*variant_index, &dest)?;
             }
 
+            Deinit(place) => {
+                let dest = self.eval_place(**place)?;
+                self.write_uninit(&dest)?;
+            }
+
             // Mark locals as alive
             StorageLive(local) => {
                 self.storage_live(*local)?;

compiler/rustc_const_eval/src/transform/check_consts/check.rs

@@ -692,6 +692,7 @@ impl<'tcx> Visitor<'tcx> for Checker<'_, 'tcx> {
         match statement.kind {
             StatementKind::Assign(..)
             | StatementKind::SetDiscriminant { .. }
+            | StatementKind::Deinit(..)
             | StatementKind::FakeRead(..)
             | StatementKind::StorageLive(_)
             | StatementKind::StorageDead(_)

compiler/rustc_const_eval/src/transform/validate.rs

@@ -346,9 +346,24 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                     self.fail(location, format!("bad arg ({:?} != usize)", op_cnt_ty))
                 }
             }
-            StatementKind::SetDiscriminant { .. } => {
-                if self.mir_phase < MirPhase::DropsLowered {
-                    self.fail(location, "`SetDiscriminant` is not allowed until drop elaboration");
+            StatementKind::SetDiscriminant { place, .. } => {
+                if self.mir_phase < MirPhase::Deaggregated {
+                    self.fail(location, "`SetDiscriminant`is not allowed until deaggregation");
+                }
+                let pty = place.ty(&self.body.local_decls, self.tcx).ty.kind();
+                if !matches!(pty, ty::Adt(..) | ty::Generator(..) | ty::Opaque(..)) {
+                    self.fail(
+                        location,
+                        format!(
+                            "`SetDiscriminant` is only allowed on ADTs and generators, not {:?}",
+                            pty
+                        ),
+                    );
+                }
+            }
+            StatementKind::Deinit(..) => {
+                if self.mir_phase < MirPhase::Deaggregated {
+                    self.fail(location, "`Deinit`is not allowed until deaggregation");
                 }
             }
             StatementKind::Retag(_, _) => {

compiler/rustc_const_eval/src/util/aggregate.rs

@@ -14,22 +14,26 @@ use std::iter::TrustedLen;
 /// (lhs as Variant).field1 = arg1;
 /// discriminant(lhs) = variant_index;  // If lhs is an enum or generator.
 pub fn expand_aggregate<'tcx>(
-    mut lhs: Place<'tcx>,
+    orig_lhs: Place<'tcx>,
     operands: impl Iterator<Item = (Operand<'tcx>, Ty<'tcx>)> + TrustedLen,
     kind: AggregateKind<'tcx>,
     source_info: SourceInfo,
     tcx: TyCtxt<'tcx>,
 ) -> impl Iterator<Item = Statement<'tcx>> + TrustedLen {
+    let mut lhs = orig_lhs;
     let mut set_discriminant = None;
     let active_field_index = match kind {
         AggregateKind::Adt(adt_did, variant_index, _, _, active_field_index) => {
             let adt_def = tcx.adt_def(adt_did);
             if adt_def.is_enum() {
                 set_discriminant = Some(Statement {
-                    kind: StatementKind::SetDiscriminant { place: Box::new(lhs), variant_index },
+                    kind: StatementKind::SetDiscriminant {
+                        place: Box::new(orig_lhs),
+                        variant_index,
+                    },
                     source_info,
                 });
-                lhs = tcx.mk_place_downcast(lhs, adt_def, variant_index);
+                lhs = tcx.mk_place_downcast(orig_lhs, adt_def, variant_index);
             }
             active_field_index
         }
@@ -38,7 +42,7 @@ pub fn expand_aggregate<'tcx>(
             // variant 0 (Unresumed).
             let variant_index = VariantIdx::new(0);
             set_discriminant = Some(Statement {
-                kind: StatementKind::SetDiscriminant { place: Box::new(lhs), variant_index },
+                kind: StatementKind::SetDiscriminant { place: Box::new(orig_lhs), variant_index },
                 source_info,
             });
 
@@ -50,27 +54,24 @@ pub fn expand_aggregate<'tcx>(
         _ => None,
     };
 
-    operands
-        .enumerate()
-        .map(move |(i, (op, ty))| {
-            let lhs_field = if let AggregateKind::Array(_) = kind {
-                let offset = u64::try_from(i).unwrap();
-                tcx.mk_place_elem(
-                    lhs,
-                    ProjectionElem::ConstantIndex {
-                        offset,
-                        min_length: offset + 1,
-                        from_end: false,
-                    },
-                )
-            } else {
-                let field = Field::new(active_field_index.unwrap_or(i));
-                tcx.mk_place_field(lhs, field, ty)
-            };
-            Statement {
-                source_info,
-                kind: StatementKind::Assign(Box::new((lhs_field, Rvalue::Use(op)))),
-            }
-        })
+    let operands = operands.enumerate().map(move |(i, (op, ty))| {
+        let lhs_field = if let AggregateKind::Array(_) = kind {
+            let offset = u64::try_from(i).unwrap();
+            tcx.mk_place_elem(
+                lhs,
+                ProjectionElem::ConstantIndex { offset, min_length: offset + 1, from_end: false },
+            )
+        } else {
+            let field = Field::new(active_field_index.unwrap_or(i));
+            tcx.mk_place_field(lhs, field, ty)
+        };
+        Statement {
+            source_info,
+            kind: StatementKind::Assign(Box::new((lhs_field, Rvalue::Use(op)))),
+        }
+    });
+    [Statement { source_info, kind: StatementKind::Deinit(Box::new(orig_lhs)) }]
+        .into_iter()
+        .chain(operands)
         .chain(set_discriminant)
 }

compiler/rustc_middle/src/mir/mod.rs

@@ -1588,8 +1588,17 @@ pub enum StatementKind<'tcx> {
     FakeRead(Box<(FakeReadCause, Place<'tcx>)>),
 
     /// Write the discriminant for a variant to the enum Place.
+    ///
+    /// This is permitted for both generators and ADTs. This does not necessarily write to the
+    /// entire place; instead, it writes to the minimum set of bytes as required by the layout for
+    /// the type.
     SetDiscriminant { place: Box<Place<'tcx>>, variant_index: VariantIdx },
 
+    /// Deinitializes the place.
+    ///
+    /// This writes `uninit` bytes to the entire place.
+    Deinit(Box<Place<'tcx>>),
+
     /// Start a live range for the storage of the local.
     StorageLive(Local),
 
@@ -1739,6 +1748,7 @@ impl Debug for Statement<'_> {
             SetDiscriminant { ref place, variant_index } => {
                 write!(fmt, "discriminant({:?}) = {:?}", place, variant_index)
             }
+            Deinit(ref place) => write!(fmt, "Deinit({:?})", place),
             AscribeUserType(box (ref place, ref c_ty), ref variance) => {
                 write!(fmt, "AscribeUserType({:?}, {:?}, {:?})", place, variance, c_ty)
             }

compiler/rustc_middle/src/mir/spanview.rs

@@ -243,6 +243,7 @@ pub fn statement_kind_name(statement: &Statement<'_>) -> &'static str {
         Assign(..) => "Assign",
         FakeRead(..) => "FakeRead",
         SetDiscriminant { .. } => "SetDiscriminant",
+        Deinit(..) => "Deinit",
         StorageLive(..) => "StorageLive",
         StorageDead(..) => "StorageDead",
         Retag(..) => "Retag",

compiler/rustc_middle/src/mir/visit.rs

@@ -395,10 +395,17 @@ macro_rules! make_mir_visitor {
                     StatementKind::SetDiscriminant { place, .. } => {
                         self.visit_place(
                             place,
-                            PlaceContext::MutatingUse(MutatingUseContext::Store),
+                            PlaceContext::MutatingUse(MutatingUseContext::SetDiscriminant),
                             location
                         );
                     }
+                    StatementKind::Deinit(place) => {
+                        self.visit_place(
+                            place,
+                            PlaceContext::MutatingUse(MutatingUseContext::Deinit),
+                            location
+                        )
+                    }
                     StatementKind::StorageLive(local) => {
                         self.visit_local(
                             local,
@@ -1174,6 +1181,10 @@ pub enum NonMutatingUseContext {
 pub enum MutatingUseContext {
     /// Appears as LHS of an assignment.
     Store,
+    /// Appears on `SetDiscriminant`
+    SetDiscriminant,
+    /// Appears on `Deinit`
+    Deinit,
     /// Output operand of an inline assembly block.
     AsmOutput,
     /// Destination of a call.