Rollup merge of #129738 - nnethercote:rustc_mir_transform-cleanups, r=cjgillot

`rustc_mir_transform` cleanups

A bunch of small improvements I made while looking closely at this code.

r? `@saethlin`

compiler/rustc_mir_transform/src/add_moves_for_packed_drops.rs

@@ -86,7 +86,7 @@ fn add_move_for_packed_drop<'tcx>(
 
     let source_info = terminator.source_info;
     let ty = place.ty(body, tcx).ty;
-    let temp = patch.new_temp(ty, terminator.source_info.span);
+    let temp = patch.new_temp(ty, source_info.span);
 
     let storage_dead_block = patch.new_block(BasicBlockData {
         statements: vec![Statement { source_info, kind: StatementKind::StorageDead(temp) }],

compiler/rustc_mir_transform/src/check_packed_ref.rs

@@ -37,24 +37,17 @@ impl<'tcx> Visitor<'tcx> for PackedRefChecker<'_, 'tcx> {
     }
 
     fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, _location: Location) {
-        if context.is_borrow() {
-            if util::is_disaligned(self.tcx, self.body, self.param_env, *place) {
-                let def_id = self.body.source.instance.def_id();
-                if let Some(impl_def_id) = self.tcx.impl_of_method(def_id)
-                    && self.tcx.is_builtin_derived(impl_def_id)
-                {
-                    // If we ever reach here it means that the generated derive
-                    // code is somehow doing an unaligned reference, which it
-                    // shouldn't do.
-                    span_bug!(
-                        self.source_info.span,
-                        "builtin derive created an unaligned reference"
-                    );
-                } else {
-                    self.tcx
-                        .dcx()
-                        .emit_err(errors::UnalignedPackedRef { span: self.source_info.span });
-                }
+        if context.is_borrow() && util::is_disaligned(self.tcx, self.body, self.param_env, *place) {
+            let def_id = self.body.source.instance.def_id();
+            if let Some(impl_def_id) = self.tcx.impl_of_method(def_id)
+                && self.tcx.is_builtin_derived(impl_def_id)
+            {
+                // If we ever reach here it means that the generated derive
+                // code is somehow doing an unaligned reference, which it
+                // shouldn't do.
+                span_bug!(self.source_info.span, "builtin derive created an unaligned reference");
+            } else {
+                self.tcx.dcx().emit_err(errors::UnalignedPackedRef { span: self.source_info.span });
             }
         }
     }

compiler/rustc_mir_transform/src/coroutine.rs

@@ -63,7 +63,9 @@ use rustc_index::bit_set::{BitMatrix, BitSet, GrowableBitSet};
 use rustc_index::{Idx, IndexVec};
 use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
 use rustc_middle::mir::*;
-use rustc_middle::ty::{self, CoroutineArgs, CoroutineArgsExt, InstanceKind, Ty, TyCtxt};
+use rustc_middle::ty::{
+    self, CoroutineArgs, CoroutineArgsExt, GenericArgsRef, InstanceKind, Ty, TyCtxt,
+};
 use rustc_middle::{bug, span_bug};
 use rustc_mir_dataflow::impls::{
     MaybeBorrowedLocals, MaybeLiveLocals, MaybeRequiresStorage, MaybeStorageLive,
@@ -113,11 +115,18 @@ impl<'tcx> MutVisitor<'tcx> for RenameLocalVisitor<'tcx> {
     }
 }
 
-struct DerefArgVisitor<'tcx> {
+struct SelfArgVisitor<'tcx> {
     tcx: TyCtxt<'tcx>,
+    new_base: Place<'tcx>,
 }
 
-impl<'tcx> MutVisitor<'tcx> for DerefArgVisitor<'tcx> {
+impl<'tcx> SelfArgVisitor<'tcx> {
+    fn new(tcx: TyCtxt<'tcx>, elem: ProjectionElem<Local, Ty<'tcx>>) -> Self {
+        Self { tcx, new_base: Place { local: SELF_ARG, projection: tcx.mk_place_elems(&[elem]) } }
+    }
+}
+
+impl<'tcx> MutVisitor<'tcx> for SelfArgVisitor<'tcx> {
     fn tcx(&self) -> TyCtxt<'tcx> {
         self.tcx
     }
@@ -128,53 +137,7 @@ impl<'tcx> MutVisitor<'tcx> for DerefArgVisitor<'tcx> {
 
     fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) {
         if place.local == SELF_ARG {
-            replace_base(
-                place,
-                Place {
-                    local: SELF_ARG,
-                    projection: self.tcx().mk_place_elems(&[ProjectionElem::Deref]),
-                },
-                self.tcx,
-            );
-        } else {
-            self.visit_local(&mut place.local, context, location);
-
-            for elem in place.projection.iter() {
-                if let PlaceElem::Index(local) = elem {
-                    assert_ne!(local, SELF_ARG);
-                }
-            }
-        }
-    }
-}
-
-struct PinArgVisitor<'tcx> {
-    ref_coroutine_ty: Ty<'tcx>,
-    tcx: TyCtxt<'tcx>,
-}
-
-impl<'tcx> MutVisitor<'tcx> for PinArgVisitor<'tcx> {
-    fn tcx(&self) -> TyCtxt<'tcx> {
-        self.tcx
-    }
-
-    fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
-        assert_ne!(*local, SELF_ARG);
-    }
-
-    fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) {
-        if place.local == SELF_ARG {
-            replace_base(
-                place,
-                Place {
-                    local: SELF_ARG,
-                    projection: self.tcx().mk_place_elems(&[ProjectionElem::Field(
-                        FieldIdx::ZERO,
-                        self.ref_coroutine_ty,
-                    )]),
-                },
-                self.tcx,
-            );
+            replace_base(place, self.new_base, self.tcx);
         } else {
             self.visit_local(&mut place.local, context, location);
 
@@ -198,15 +161,6 @@ fn replace_base<'tcx>(place: &mut Place<'tcx>, new_base: Place<'tcx>, tcx: TyCtx
 
 const SELF_ARG: Local = Local::from_u32(1);
 
-/// Coroutine has not been resumed yet.
-const UNRESUMED: usize = CoroutineArgs::UNRESUMED;
-/// Coroutine has returned / is completed.
-const RETURNED: usize = CoroutineArgs::RETURNED;
-/// Coroutine has panicked and is poisoned.
-const POISONED: usize = CoroutineArgs::POISONED;
-/// Number of reserved variants of coroutine state.
-const RESERVED_VARIANTS: usize = CoroutineArgs::RESERVED_VARIANTS;
-
 /// A `yield` point in the coroutine.
 struct SuspensionPoint<'tcx> {
     /// State discriminant used when suspending or resuming at this point.
@@ -261,14 +215,10 @@ impl<'tcx> TransformVisitor<'tcx> {
             // `gen` continues return `None`
             CoroutineKind::Desugared(CoroutineDesugaring::Gen, _) => {
                 let option_def_id = self.tcx.require_lang_item(LangItem::Option, None);
-                Rvalue::Aggregate(
-                    Box::new(AggregateKind::Adt(
-                        option_def_id,
-                        VariantIdx::ZERO,
-                        self.tcx.mk_args(&[self.old_yield_ty.into()]),
-                        None,
-                        None,
-                    )),
+                make_aggregate_adt(
+                    option_def_id,
+                    VariantIdx::ZERO,
+                    self.tcx.mk_args(&[self.old_yield_ty.into()]),
                     IndexVec::new(),
                 )
             }
@@ -317,64 +267,28 @@ impl<'tcx> TransformVisitor<'tcx> {
         is_return: bool,
         statements: &mut Vec<Statement<'tcx>>,
     ) {
+        const ZERO: VariantIdx = VariantIdx::ZERO;
+        const ONE: VariantIdx = VariantIdx::from_usize(1);
         let rvalue = match self.coroutine_kind {
             CoroutineKind::Desugared(CoroutineDesugaring::Async, _) => {
                 let poll_def_id = self.tcx.require_lang_item(LangItem::Poll, None);
                 let args = self.tcx.mk_args(&[self.old_ret_ty.into()]);
-                if is_return {
-                    // Poll::Ready(val)
-                    Rvalue::Aggregate(
-                        Box::new(AggregateKind::Adt(
-                            poll_def_id,
-                            VariantIdx::ZERO,
-                            args,
-                            None,
-                            None,
-                        )),
-                        IndexVec::from_raw(vec![val]),
-                    )
+                let (variant_idx, operands) = if is_return {
+                    (ZERO, IndexVec::from_raw(vec![val])) // Poll::Ready(val)
                 } else {
-                    // Poll::Pending
-                    Rvalue::Aggregate(
-                        Box::new(AggregateKind::Adt(
-                            poll_def_id,
-                            VariantIdx::from_usize(1),
-                            args,
-                            None,
-                            None,
-                        )),
-                        IndexVec::new(),
-                    )
-                }
+                    (ONE, IndexVec::new()) // Poll::Pending
+                };
+                make_aggregate_adt(poll_def_id, variant_idx, args, operands)
             }
             CoroutineKind::Desugared(CoroutineDesugaring::Gen, _) => {
                 let option_def_id = self.tcx.require_lang_item(LangItem::Option, None);
                 let args = self.tcx.mk_args(&[self.old_yield_ty.into()]);
-                if is_return {
-                    // None
-                    Rvalue::Aggregate(
-                        Box::new(AggregateKind::Adt(
-                            option_def_id,
-                            VariantIdx::ZERO,
-                            args,
-                            None,
-                            None,
-                        )),
-                        IndexVec::new(),
-                    )
+                let (variant_idx, operands) = if is_return {
+                    (ZERO, IndexVec::new()) // None
                 } else {
-                    // Some(val)
-                    Rvalue::Aggregate(
-                        Box::new(AggregateKind::Adt(
-                            option_def_id,
-                            VariantIdx::from_usize(1),
-                            args,
-                            None,
-                            None,
-                        )),
-                        IndexVec::from_raw(vec![val]),
-                    )
-                }
+                    (ONE, IndexVec::from_raw(vec![val])) // Some(val)
+                };
+                make_aggregate_adt(option_def_id, variant_idx, args, operands)
             }
             CoroutineKind::Desugared(CoroutineDesugaring::AsyncGen, _) => {
                 if is_return {
@@ -400,31 +314,17 @@ impl<'tcx> TransformVisitor<'tcx> {
                 let coroutine_state_def_id =
                     self.tcx.require_lang_item(LangItem::CoroutineState, None);
                 let args = self.tcx.mk_args(&[self.old_yield_ty.into(), self.old_ret_ty.into()]);
-                if is_return {
-                    // CoroutineState::Complete(val)
-                    Rvalue::Aggregate(
-                        Box::new(AggregateKind::Adt(
-                            coroutine_state_def_id,
-                            VariantIdx::from_usize(1),
-                            args,
-                            None,
-                            None,
-                        )),
-                        IndexVec::from_raw(vec![val]),
-                    )
+                let variant_idx = if is_return {
+                    ONE // CoroutineState::Complete(val)
                 } else {
-                    // CoroutineState::Yielded(val)
-                    Rvalue::Aggregate(
-                        Box::new(AggregateKind::Adt(
-                            coroutine_state_def_id,
-                            VariantIdx::ZERO,
-                            args,
-                            None,
-                            None,
-                        )),
-                        IndexVec::from_raw(vec![val]),
-                    )
-                }
+                    ZERO // CoroutineState::Yielded(val)
+                };
+                make_aggregate_adt(
+                    coroutine_state_def_id,
+                    variant_idx,
+                    args,
+                    IndexVec::from_raw(vec![val]),
+                )
             }
         };
 
@@ -517,7 +417,7 @@ impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> {
             self.make_state(v, source_info, is_return, &mut data.statements);
             let state = if let Some((resume, mut resume_arg)) = resume {
                 // Yield
-                let state = RESERVED_VARIANTS + self.suspension_points.len();
+                let state = CoroutineArgs::RESERVED_VARIANTS + self.suspension_points.len();
 
                 // The resume arg target location might itself be remapped if its base local is
                 // live across a yield.
@@ -550,7 +450,7 @@ impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> {
                 VariantIdx::new(state)
             } else {
                 // Return
-                VariantIdx::new(RETURNED) // state for returned
+                VariantIdx::new(CoroutineArgs::RETURNED) // state for returned
             };
             data.statements.push(self.set_discr(state, source_info));
             data.terminator_mut().kind = TerminatorKind::Return;
@@ -560,6 +460,15 @@ impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> {
     }
 }
 
+fn make_aggregate_adt<'tcx>(
+    def_id: DefId,
+    variant_idx: VariantIdx,
+    args: GenericArgsRef<'tcx>,
+    operands: IndexVec<FieldIdx, Operand<'tcx>>,
+) -> Rvalue<'tcx> {
+    Rvalue::Aggregate(Box::new(AggregateKind::Adt(def_id, variant_idx, args, None, None)), operands)
+}
+
 fn make_coroutine_state_argument_indirect<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     let coroutine_ty = body.local_decls.raw[1].ty;
 
@@ -569,7 +478,7 @@ fn make_coroutine_state_argument_indirect<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Bo
     body.local_decls.raw[1].ty = ref_coroutine_ty;
 
     // Add a deref to accesses of the coroutine state
-    DerefArgVisitor { tcx }.visit_body(body);
+    SelfArgVisitor::new(tcx, ProjectionElem::Deref).visit_body(body);
 }
 
 fn make_coroutine_state_argument_pinned<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
@@ -584,7 +493,8 @@ fn make_coroutine_state_argument_pinned<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body
     body.local_decls.raw[1].ty = pin_ref_coroutine_ty;
 
     // Add the Pin field access to accesses of the coroutine state
-    PinArgVisitor { ref_coroutine_ty, tcx }.visit_body(body);
+    SelfArgVisitor::new(tcx, ProjectionElem::Field(FieldIdx::ZERO, ref_coroutine_ty))
+        .visit_body(body);
 }
 
 /// Allocates a new local and replaces all references of `local` with it. Returns the new local.
@@ -651,8 +561,6 @@ fn transform_async_context<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
                         let local = eliminate_get_context_call(&mut body[bb]);
                         replace_resume_ty_local(tcx, body, local, context_mut_ref);
                     }
-                } else {
-                    continue;
                 }
             }
             TerminatorKind::Yield { resume_arg, .. } => {
@@ -665,24 +573,23 @@ fn transform_async_context<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
 
 fn eliminate_get_context_call<'tcx>(bb_data: &mut BasicBlockData<'tcx>) -> Local {
     let terminator = bb_data.terminator.take().unwrap();
-    if let TerminatorKind::Call { args, destination, target, .. } = terminator.kind {
-        let [arg] = *Box::try_from(args).unwrap();
-        let local = arg.node.place().unwrap().local;
-
-        let arg = Rvalue::Use(arg.node);
-        let assign = Statement {
-            source_info: terminator.source_info,
-            kind: StatementKind::Assign(Box::new((destination, arg))),
-        };
-        bb_data.statements.push(assign);
-        bb_data.terminator = Some(Terminator {
-            source_info: terminator.source_info,
-            kind: TerminatorKind::Goto { target: target.unwrap() },
-        });
-        local
-    } else {
+    let TerminatorKind::Call { args, destination, target, .. } = terminator.kind else {
         bug!();
-    }
+    };
+    let [arg] = *Box::try_from(args).unwrap();
+    let local = arg.node.place().unwrap().local;
+
+    let arg = Rvalue::Use(arg.node);
+    let assign = Statement {
+        source_info: terminator.source_info,
+        kind: StatementKind::Assign(Box::new((destination, arg))),
+    };
+    bb_data.statements.push(assign);
+    bb_data.terminator = Some(Terminator {
+        source_info: terminator.source_info,
+        kind: TerminatorKind::Goto { target: target.unwrap() },
+    });
+    local
 }
 
 #[cfg_attr(not(debug_assertions), allow(unused))]
@@ -1085,10 +992,11 @@ fn compute_layout<'tcx>(
     // Build the coroutine variant field list.
     // Create a map from local indices to coroutine struct indices.
     let mut variant_fields: IndexVec<VariantIdx, IndexVec<FieldIdx, CoroutineSavedLocal>> =
-        iter::repeat(IndexVec::new()).take(RESERVED_VARIANTS).collect();
+        iter::repeat(IndexVec::new()).take(CoroutineArgs::RESERVED_VARIANTS).collect();
     let mut remap = IndexVec::from_elem_n(None, saved_locals.domain_size());
     for (suspension_point_idx, live_locals) in live_locals_at_suspension_points.iter().enumerate() {
-        let variant_index = VariantIdx::from(RESERVED_VARIANTS + suspension_point_idx);
+        let variant_index =
+            VariantIdx::from(CoroutineArgs::RESERVED_VARIANTS + suspension_point_idx);
         let mut fields = IndexVec::new();
         for (idx, saved_local) in live_locals.iter().enumerate() {
             fields.push(saved_local);
@@ -1183,12 +1091,10 @@ fn elaborate_coroutine_drops<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
                 source_info,
                 kind: TerminatorKind::Drop { place, target, unwind, replace: _ },
             } => {
-                if let Some(local) = place.as_local() {
-                    if local == SELF_ARG {
-                        (target, unwind, source_info)
-                    } else {
-                        continue;
-                    }
+                if let Some(local) = place.as_local()
+                    && local == SELF_ARG
+                {
+                    (target, unwind, source_info)
                 } else {
                     continue;
                 }
@@ -1237,7 +1143,7 @@ fn create_coroutine_drop_shim<'tcx>(
 
     let mut cases = create_cases(&mut body, transform, Operation::Drop);
 
-    cases.insert(0, (UNRESUMED, drop_clean));
+    cases.insert(0, (CoroutineArgs::UNRESUMED, drop_clean));
 
     // The returned state and the poisoned state fall through to the default
     // case which is just to return
@@ -1387,7 +1293,9 @@ fn create_coroutine_resume_function<'tcx>(
     if can_unwind {
         let source_info = SourceInfo::outermost(body.span);
         let poison_block = body.basic_blocks_mut().push(BasicBlockData {
-            statements: vec![transform.set_discr(VariantIdx::new(POISONED), source_info)],
+            statements: vec![
+                transform.set_discr(VariantIdx::new(CoroutineArgs::POISONED), source_info),
+            ],
             terminator: Some(Terminator { source_info, kind: TerminatorKind::UnwindResume }),
             is_cleanup: true,
         });
@@ -1419,13 +1327,16 @@ fn create_coroutine_resume_function<'tcx>(
     use rustc_middle::mir::AssertKind::{ResumedAfterPanic, ResumedAfterReturn};
 
     // Jump to the entry point on the unresumed
-    cases.insert(0, (UNRESUMED, START_BLOCK));
+    cases.insert(0, (CoroutineArgs::UNRESUMED, START_BLOCK));
 
     // Panic when resumed on the returned or poisoned state
     if can_unwind {
         cases.insert(
             1,
-            (POISONED, insert_panic_block(tcx, body, ResumedAfterPanic(transform.coroutine_kind))),
+            (
+                CoroutineArgs::POISONED,
+                insert_panic_block(tcx, body, ResumedAfterPanic(transform.coroutine_kind)),
+            ),
         );
     }
 
@@ -1440,7 +1351,7 @@ fn create_coroutine_resume_function<'tcx>(
                 transform.insert_none_ret_block(body)
             }
         };
-        cases.insert(1, (RETURNED, block));
+        cases.insert(1, (CoroutineArgs::RETURNED, block));
     }
 
     insert_switch(body, cases, &transform, TerminatorKind::Unreachable);

compiler/rustc_mir_transform/src/lib.rs

@@ -3,13 +3,11 @@
 #![feature(box_patterns)]
 #![feature(const_type_name)]
 #![feature(cow_is_borrowed)]
-#![feature(decl_macro)]
 #![feature(if_let_guard)]
 #![feature(impl_trait_in_assoc_type)]
 #![feature(let_chains)]
 #![feature(map_try_insert)]
 #![feature(never_type)]
-#![feature(option_get_or_insert_default)]
 #![feature(round_char_boundary)]
 #![feature(try_blocks)]
 #![feature(yeet_expr)]
@@ -17,6 +15,7 @@
 
 use hir::ConstContext;
 use required_consts::RequiredConstsVisitor;
+use rustc_const_eval::check_consts::{self, ConstCx};
 use rustc_const_eval::util;
 use rustc_data_structures::fx::FxIndexSet;
 use rustc_data_structures::steal::Steal;
@@ -33,6 +32,7 @@ use rustc_middle::mir::{
 use rustc_middle::ty::{self, TyCtxt, TypeVisitableExt};
 use rustc_middle::util::Providers;
 use rustc_middle::{bug, query, span_bug};
+use rustc_mir_dataflow::rustc_peek;
 use rustc_span::source_map::Spanned;
 use rustc_span::{sym, DUMMY_SP};
 use rustc_trait_selection::traits;
@@ -108,9 +108,6 @@ mod unreachable_enum_branching;
 mod unreachable_prop;
 mod validate;
 
-use rustc_const_eval::check_consts::{self, ConstCx};
-use rustc_mir_dataflow::rustc_peek;
-
 rustc_fluent_macro::fluent_messages! { "../messages.ftl" }
 
 pub fn provide(providers: &mut Providers) {
@@ -128,7 +125,7 @@ pub fn provide(providers: &mut Providers) {
         mir_coroutine_witnesses: coroutine::mir_coroutine_witnesses,
         optimized_mir,
         is_mir_available,
-        is_ctfe_mir_available: |tcx, did| is_mir_available(tcx, did),
+        is_ctfe_mir_available: is_mir_available,
         mir_callgraph_reachable: inline::cycle::mir_callgraph_reachable,
         mir_inliner_callees: inline::cycle::mir_inliner_callees,
         promoted_mir,
@@ -221,10 +218,8 @@ fn is_mir_available(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
 /// Finds the full set of `DefId`s within the current crate that have
 /// MIR associated with them.
 fn mir_keys(tcx: TyCtxt<'_>, (): ()) -> FxIndexSet<LocalDefId> {
-    let mut set = FxIndexSet::default();
-
     // All body-owners have MIR associated with them.
-    set.extend(tcx.hir().body_owners());
+    let mut set: FxIndexSet<_> = tcx.hir().body_owners().collect();
 
     // Additionally, tuple struct/variant constructors have MIR, but
     // they don't have a BodyId, so we need to build them separately.
@@ -249,8 +244,7 @@ fn mir_const_qualif(tcx: TyCtxt<'_>, def: LocalDefId) -> ConstQualifs {
 
     // No need to const-check a non-const `fn`.
     match const_kind {
-        Some(ConstContext::Const { .. } | ConstContext::Static(_))
-        | Some(ConstContext::ConstFn) => {}
+        Some(ConstContext::Const { .. } | ConstContext::Static(_) | ConstContext::ConstFn) => {}
         None => span_bug!(
             tcx.def_span(def),
             "`mir_const_qualif` should only be called on const fns and const items"