Rollup merge of #77568 - lcnr:mir-inline-def-id, r=ecstatic-morse

inliner: use caller param_env

We used the callee param env instead of the caller param env by accident in #77430, this PR fixes that and caches it in the `Inliner` struct.

fixes #77564

r? @ecstatic-morse

Author: Dylan-DPC

compiler/rustc_mir/src/transform/inline.rs

@@ -45,15 +45,20 @@ impl<'tcx> MirPass<'tcx> for Inline {
                 // based function.
                 debug!("function inlining is disabled when compiling with `instrument_coverage`");
             } else {
-                Inliner { tcx, codegen_fn_attrs: tcx.codegen_fn_attrs(body.source.def_id()) }
-                    .run_pass(body);
+                Inliner {
+                    tcx,
+                    param_env: tcx.param_env_reveal_all_normalized(body.source.def_id()),
+                    codegen_fn_attrs: tcx.codegen_fn_attrs(body.source.def_id()),
+                }
+                .run_pass(body);
             }
         }
     }
 }
 
 struct Inliner<'tcx> {
     tcx: TyCtxt<'tcx>,
+    param_env: ParamEnv<'tcx>,
     codegen_fn_attrs: &'tcx CodegenFnAttrs,
 }
 
@@ -75,17 +80,13 @@ impl Inliner<'tcx> {
 
         let def_id = caller_body.source.def_id();
 
-        let param_env = self.tcx.param_env_reveal_all_normalized(def_id);
-
         // Only do inlining into fn bodies.
         let self_hir_id = self.tcx.hir().local_def_id_to_hir_id(def_id.expect_local());
         if self.tcx.hir().body_owner_kind(self_hir_id).is_fn_or_closure()
             && caller_body.source.promoted.is_none()
         {
             for (bb, bb_data) in caller_body.basic_blocks().iter_enumerated() {
-                if let Some(callsite) =
-                    self.get_valid_function_call(bb, bb_data, caller_body, param_env)
-                {
+                if let Some(callsite) = self.get_valid_function_call(bb, bb_data, caller_body) {
                     callsites.push_back(callsite);
                 }
             }
@@ -131,7 +132,7 @@ impl Inliner<'tcx> {
                 let callee_body = if self.consider_optimizing(callsite, callee_body) {
                     self.tcx.subst_and_normalize_erasing_regions(
                         &callsite.substs,
-                        param_env,
+                        self.param_env,
                         callee_body,
                     )
                 } else {
@@ -159,7 +160,7 @@ impl Inliner<'tcx> {
                 // Add callsites from inlined function
                 for (bb, bb_data) in caller_body.basic_blocks().iter_enumerated().skip(start) {
                     if let Some(new_callsite) =
-                        self.get_valid_function_call(bb, bb_data, caller_body, param_env)
+                        self.get_valid_function_call(bb, bb_data, caller_body)
                     {
                         // Don't inline the same function multiple times.
                         if callsite.callee != new_callsite.callee {
@@ -190,7 +191,6 @@ impl Inliner<'tcx> {
         bb: BasicBlock,
         bb_data: &BasicBlockData<'tcx>,
         caller_body: &Body<'tcx>,
-        param_env: ParamEnv<'tcx>,
     ) -> Option<CallSite<'tcx>> {
         // Don't inline calls that are in cleanup blocks.
         if bb_data.is_cleanup {
@@ -201,8 +201,9 @@ impl Inliner<'tcx> {
         let terminator = bb_data.terminator();
         if let TerminatorKind::Call { func: ref op, .. } = terminator.kind {
             if let ty::FnDef(callee_def_id, substs) = *op.ty(caller_body, self.tcx).kind() {
-                let instance =
-                    Instance::resolve(self.tcx, param_env, callee_def_id, substs).ok().flatten()?;
+                let instance = Instance::resolve(self.tcx, self.param_env, callee_def_id, substs)
+                    .ok()
+                    .flatten()?;
 
                 if let InstanceDef::Virtual(..) = instance.def {
                     return None;
@@ -300,9 +301,6 @@ impl Inliner<'tcx> {
         debug!("    final inline threshold = {}", threshold);
 
         // FIXME: Give a bonus to functions with only a single caller
-
-        let param_env = tcx.param_env(callee_body.source.def_id());
-
         let mut first_block = true;
         let mut cost = 0;
 
@@ -335,7 +333,7 @@ impl Inliner<'tcx> {
                     // If the place doesn't actually need dropping, treat it like
                     // a regular goto.
                     let ty = place.ty(callee_body, tcx).subst(tcx, callsite.substs).ty;
-                    if ty.needs_drop(tcx, param_env) {
+                    if ty.needs_drop(tcx, self.param_env) {
                         cost += CALL_PENALTY;
                         if let Some(unwind) = unwind {
                             cost += LANDINGPAD_PENALTY;
@@ -400,7 +398,7 @@ impl Inliner<'tcx> {
             let ty = v.ty.subst(tcx, callsite.substs);
             // Cost of the var is the size in machine-words, if we know
             // it.
-            if let Some(size) = type_size_of(tcx, param_env, ty) {
+            if let Some(size) = type_size_of(tcx, self.param_env, ty) {
                 cost += (size / ptr_size) as usize;
             } else {
                 cost += UNKNOWN_SIZE_COST;

src/test/ui/mir/mir-inlining/ice-issue-77564.rs

@@ -0,0 +1,38 @@
+// run-pass
+// compile-flags:-Zmir-opt-level=2
+
+use std::mem::MaybeUninit;
+const N: usize = 2;
+
+trait CollectArray<A>: Iterator<Item = A> {
+    fn inner_array(&mut self) -> [A; N];
+    fn collect_array(&mut self) -> [A; N] {
+        let result = self.inner_array();
+        assert!(self.next().is_none());
+        result
+    }
+}
+
+impl<A, I: ?Sized> CollectArray<A> for I
+where
+    I: Iterator<Item = A>,
+{
+    fn inner_array(&mut self) -> [A; N] {
+        let mut result: [MaybeUninit<A>; N] = unsafe { MaybeUninit::uninit().assume_init() };
+        for (dest, item) in result.iter_mut().zip(self) {
+            *dest = MaybeUninit::new(item);
+        }
+        let temp_ptr: *const [MaybeUninit<A>; N] = &result;
+        unsafe { std::ptr::read(temp_ptr as *const [A; N]) }
+    }
+}
+
+fn main() {
+    assert_eq!(
+        [[1, 2], [3, 4]]
+            .iter()
+            .map(|row| row.iter().collect_array())
+            .collect_array(),
+        [[&1, &2], [&3, &4]]
+    );
+}