Auto merge of #145477 - cjgillot:codegen-mode, r=oli-obk

Introduce `TypingMode::Codegen` to avoid layout cycles on coroutines



Computing layout of coroutines depends on their `optimized_mir`. At the same time, MIR opts can require using layouts to work. For instance to evaluate constants. This leads to cycles and clumsy workarounds.

This PR creates a new typing mode for layout computations:
- when a coroutine's layout is requested with `TypingMode::PostAnalysis` or earlier, return `LayourError::TooGeneric`;
- when a coroutine's layout is requested with `TypingMode::Codegen`, actually compute it.

`TypingMode::Codegen` is meant be be used by codegen code, and analyses that require coroutine layout, like transmute check and coroutine recursion check.

With this PR, we can remove all `is_coroutine` checks from `rustc_mir_transform` and unlock simplifying coroutine MIR.

Perf is not terrific. This PR causes recomputation of a few queries, and I had to insert workarounds.

Author: bors

compiler/rustc_const_eval/src/const_eval/eval_queries.rs

@@ -8,7 +8,7 @@ use rustc_middle::mir::{self, ConstAlloc, ConstValue};
 use rustc_middle::query::TyCtxtAt;
 use rustc_middle::ty::layout::{HasTypingEnv, TyAndLayout};
 use rustc_middle::ty::print::with_no_trimmed_paths;
-use rustc_middle::ty::{self, Ty, TyCtxt};
+use rustc_middle::ty::{self, Ty, TyCtxt, TypeVisitable};
 use rustc_middle::{bug, throw_inval};
 use rustc_span::Span;
 use rustc_span::def_id::LocalDefId;
@@ -23,6 +23,33 @@ use crate::interpret::{
 };
 use crate::{CTRL_C_RECEIVED, errors};
 
+fn retry_codegen_mode_with_postanalysis<'tcx, K: TypeVisitable<TyCtxt<'tcx>>, V>(
+    key: ty::PseudoCanonicalInput<'tcx, K>,
+    f: impl FnOnce(ty::PseudoCanonicalInput<'tcx, K>) -> Result<V, ErrorHandled>,
+) -> Option<Result<V, ErrorHandled>> {
+    let ty::PseudoCanonicalInput { typing_env, value } = key;
+    match typing_env.typing_mode().assert_not_erased() {
+        // We are in codegen. It's very likely this constant has been evaluated in PostAnalysis
+        // before. Try to reuse this evaluation, and only re-run if we hit a `TooGeneric` error.
+        ty::TypingMode::Codegen => {
+            let with_postanalysis =
+                ty::TypingEnv::new(typing_env.param_env, ty::TypingMode::PostAnalysis);
+            let with_postanalysis = f(with_postanalysis.as_query_input(value));
+            match with_postanalysis {
+                Ok(_) | Err(ErrorHandled::Reported(..)) => return Some(with_postanalysis),
+                Err(ErrorHandled::TooGeneric(_)) => {}
+            }
+        }
+        ty::TypingMode::Coherence
+        | ty::TypingMode::Analysis { .. }
+        | ty::TypingMode::Borrowck { .. }
+        | ty::TypingMode::PostBorrowckAnalysis { .. }
+        | ty::TypingMode::PostAnalysis => {}
+    }
+
+    None
+}
+
 fn setup_for_eval<'tcx>(
     ecx: &mut CompileTimeInterpCx<'tcx>,
     cid: GlobalId<'tcx>,
@@ -327,9 +354,18 @@ pub fn eval_to_const_value_raw_provider<'tcx>(
     tcx: TyCtxt<'tcx>,
     key: ty::PseudoCanonicalInput<'tcx, GlobalId<'tcx>>,
 ) -> ::rustc_middle::mir::interpret::EvalToConstValueResult<'tcx> {
+    crate::assert_typing_mode(key.typing_env.typing_mode());
+
     if let Some((value, _ty)) = tcx.trivial_const(key.value.instance.def_id()) {
         return Ok(value);
     }
+
+    if let Some(retry) =
+        retry_codegen_mode_with_postanalysis(key, |key| tcx.eval_to_const_value_raw(key))
+    {
+        return retry;
+    }
+
     tcx.eval_to_allocation_raw(key).map(|val| turn_into_const_value(tcx, val, key))
 }
 
@@ -369,23 +405,18 @@ pub fn eval_to_allocation_raw_provider<'tcx>(
     tcx: TyCtxt<'tcx>,
     key: ty::PseudoCanonicalInput<'tcx, GlobalId<'tcx>>,
 ) -> ::rustc_middle::mir::interpret::EvalToAllocationRawResult<'tcx> {
+    crate::assert_typing_mode(key.typing_env.typing_mode());
+    if let Some(retry) =
+        retry_codegen_mode_with_postanalysis(key, |key| tcx.eval_to_allocation_raw(key))
+    {
+        return retry;
+    }
+
     // This shouldn't be used for statics, since statics are conceptually places,
     // not values -- so what we do here could break pointer identity.
     assert!(key.value.promoted.is_some() || !tcx.is_static(key.value.instance.def_id()));
 
     if cfg!(debug_assertions) {
-        match key.typing_env.typing_mode().assert_not_erased() {
-            ty::TypingMode::PostAnalysis => {}
-            ty::TypingMode::Coherence
-            | ty::TypingMode::Analysis { .. }
-            | ty::TypingMode::Borrowck { .. }
-            | ty::TypingMode::PostBorrowckAnalysis { .. } => {
-                bug!(
-                    "Const eval should always happens in PostAnalysis mode. See the comment in `InterpCx::new` for more details."
-                )
-            }
-        }
-
         // Make sure we format the instance even if we do not print it.
         // This serves as a regression test against an ICE on printing.
         // The next two lines concatenated contain some discussion:

compiler/rustc_const_eval/src/const_eval/valtrees.rs

@@ -241,19 +241,7 @@ pub(crate) fn eval_to_valtree<'tcx>(
     typing_env: ty::TypingEnv<'tcx>,
     cid: GlobalId<'tcx>,
 ) -> EvalToValTreeResult<'tcx> {
-    if cfg!(debug_assertions) {
-        match typing_env.typing_mode().assert_not_erased() {
-            ty::TypingMode::PostAnalysis => {}
-            ty::TypingMode::Coherence
-            | ty::TypingMode::Analysis { .. }
-            | ty::TypingMode::Borrowck { .. }
-            | ty::TypingMode::PostBorrowckAnalysis { .. } => {
-                bug!(
-                    "Const eval should always happens in PostAnalysis mode. See the comment in `InterpCx::new` for more details."
-                )
-            }
-        }
-    }
+    crate::assert_typing_mode(typing_env.typing_mode());
     let const_alloc = tcx.eval_to_allocation_raw(typing_env.as_query_input(cid))?;
 
     // FIXME Need to provide a span to `eval_to_valtree`

compiler/rustc_const_eval/src/interpret/eval_context.rs

@@ -9,10 +9,9 @@ use rustc_middle::ty::layout::{
     LayoutOfHelpers, TyAndLayout,
 };
 use rustc_middle::ty::{
-    self, GenericArgsRef, Ty, TyCtxt, TypeFoldable, TypeVisitableExt, TypingEnv, TypingMode,
-    Variance,
+    self, GenericArgsRef, Ty, TyCtxt, TypeFoldable, TypeVisitableExt, TypingEnv, Variance,
 };
-use rustc_middle::{bug, mir, span_bug};
+use rustc_middle::{mir, span_bug};
 use rustc_span::Span;
 use rustc_target::callconv::FnAbi;
 use tracing::{debug, trace};
@@ -243,21 +242,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         typing_env: ty::TypingEnv<'tcx>,
         machine: M,
     ) -> Self {
-        // Const eval always happens in post analysis mode in order to be able to use the hidden types of
-        // opaque types. This is needed for trivial things like `size_of`, but also for using associated
-        // types that are not specified in the opaque type. We also use MIR bodies whose opaque types have
-        // already been revealed, so we'd be able to at least partially observe the hidden types anyways.
-        if cfg!(debug_assertions) {
-            match typing_env.typing_mode().assert_not_erased() {
-                TypingMode::PostAnalysis => {}
-                TypingMode::Coherence
-                | TypingMode::Analysis { .. }
-                | TypingMode::Borrowck { .. }
-                | TypingMode::PostBorrowckAnalysis { .. } => {
-                    bug!("Const eval should always happens in PostAnalysis mode.");
-                }
-            }
-        }
+        crate::assert_typing_mode(typing_env.typing_mode());
 
         InterpCx {
             machine,

compiler/rustc_const_eval/src/lib.rs

@@ -18,8 +18,28 @@ pub mod util;
 
 use std::sync::atomic::AtomicBool;
 
-use rustc_middle::ty;
 use rustc_middle::util::Providers;
+use rustc_middle::{bug, ty};
+
+/// Const eval always happens in post analysis mode in order to be able to use the hidden types of
+/// opaque types. This is needed for trivial things like `size_of`, but also for using associated
+/// types that are not specified in the opaque type. We also use MIR bodies whose opaque types have
+/// already been revealed, so we'd be able to at least partially observe the hidden types anyways.
+fn assert_typing_mode(typing_mode: ty::TypingMode<'_>) {
+    if cfg!(debug_assertions) {
+        match typing_mode.assert_not_erased() {
+            ty::TypingMode::PostAnalysis | ty::TypingMode::Codegen => {}
+            // Const eval always happens in PostAnalysis or Codegen mode. See the comment in
+            // `InterpCx::new` for more details.
+            ty::TypingMode::Coherence
+            | ty::TypingMode::Analysis { .. }
+            | ty::TypingMode::Borrowck { .. }
+            | ty::TypingMode::PostBorrowckAnalysis { .. } => bug!(
+                "Const eval should always happens in PostAnalysis or Codegen mode. See the comment on `assert_typing_mode` for more details."
+            ),
+        }
+    }
+}
 
 pub fn provide(providers: &mut Providers) {
     const_eval::provide(&mut providers.queries);

compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs

@@ -663,7 +663,8 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             ty::TypingMode::Coherence
             | ty::TypingMode::Borrowck { .. }
             | ty::TypingMode::PostBorrowckAnalysis { .. }
-            | ty::TypingMode::PostAnalysis => {
+            | ty::TypingMode::PostAnalysis
+            | ty::TypingMode::Codegen => {
                 bug!()
             }
         };

compiler/rustc_hir_typeck/src/intrinsicck.rs

@@ -147,7 +147,7 @@ pub(crate) fn check_transmutes(tcx: TyCtxt<'_>, owner: LocalDefId) -> Result<(),
         return Err(e);
     };
 
-    let typing_env = ty::TypingEnv::post_analysis(tcx, owner);
+    let typing_env = ty::TypingEnv::codegen(tcx, owner);
     let mut result = Ok(());
     for &(from, to, hir_id) in &typeck_results.transmutes_to_check {
         result = result.and(check_transmute(tcx, typing_env, from, to, hir_id));

compiler/rustc_hir_typeck/src/opaque_types.rs

@@ -105,7 +105,8 @@ impl<'tcx> FnCtxt<'_, 'tcx> {
             ty::TypingMode::Coherence
             | ty::TypingMode::Borrowck { .. }
             | ty::TypingMode::PostBorrowckAnalysis { .. }
-            | ty::TypingMode::PostAnalysis => {
+            | ty::TypingMode::PostAnalysis
+            | ty::TypingMode::Codegen => {
                 bug!()
             }
         };

compiler/rustc_infer/src/infer/mod.rs

@@ -350,7 +350,8 @@ impl<'tcx> Drop for InferCtxt<'tcx> {
             TypingMode::Coherence
             | TypingMode::Analysis { .. }
             | TypingMode::PostBorrowckAnalysis { .. }
-            | TypingMode::PostAnalysis => {}
+            | TypingMode::PostAnalysis
+            | TypingMode::Codegen => {}
             // In erased mode, the opaque type storage is always empty
             TypingMode::ErasedNotCoherence(..) => {}
             TypingMode::Borrowck { .. } => {
@@ -1156,7 +1157,8 @@ impl<'tcx> InferCtxt<'tcx> {
             // to support PostBorrowckAnalysis in the old solver as well.
             TypingMode::Coherence
             | TypingMode::PostBorrowckAnalysis { .. }
-            | TypingMode::PostAnalysis => false,
+            | TypingMode::PostAnalysis
+            | TypingMode::Codegen => false,
         }
     }
 
@@ -1486,7 +1488,8 @@ impl<'tcx> InferCtxt<'tcx> {
             }
             mode @ (ty::TypingMode::Coherence
             | ty::TypingMode::PostBorrowckAnalysis { .. }
-            | ty::TypingMode::PostAnalysis) => mode,
+            | ty::TypingMode::PostAnalysis
+            | ty::TypingMode::Codegen) => mode,
             ty::TypingMode::ErasedNotCoherence(MayBeErased) => unreachable!(),
         };
         ty::TypingEnv::new(param_env, typing_mode)

compiler/rustc_infer/src/infer/opaque_types/mod.rs

@@ -283,7 +283,9 @@ impl<'tcx> InferCtxt<'tcx> {
                         .map(|obligation| obligation.as_goal()),
                 );
             }
-            mode @ (ty::TypingMode::PostBorrowckAnalysis { .. } | ty::TypingMode::PostAnalysis) => {
+            mode @ (ty::TypingMode::PostBorrowckAnalysis { .. }
+            | ty::TypingMode::PostAnalysis
+            | ty::TypingMode::Codegen) => {
                 bug!("insert hidden type in {mode:?}")
             }
         }

compiler/rustc_interface/src/passes.rs

@@ -1180,10 +1180,10 @@ fn run_required_analyses(tcx: TyCtxt<'_>) {
                 && (!tcx.is_async_drop_in_place_coroutine(def_id.to_def_id()))
             {
                 // Eagerly check the unsubstituted layout for cycles.
-                tcx.ensure_ok().layout_of(
-                    ty::TypingEnv::post_analysis(tcx, def_id.to_def_id())
-                        .as_query_input(tcx.type_of(def_id).instantiate_identity().skip_norm_wip()),
-                );
+                tcx.ensure_ok()
+                    .layout_of(ty::TypingEnv::codegen(tcx, def_id.to_def_id()).as_query_input(
+                        tcx.type_of(def_id).instantiate_identity().skip_norm_wip(),
+                    ));
             }
         });
     });