Store the body type in THIR.

4 files changed, 57 insertions, 37 deletions

diff --git a/compiler/rustc_middle/src/thir.rs b/compiler/rustc_middle/src/thir.rs
index 6f2dac46753..f9820fdd394 100644
--- a/compiler/rustc_middle/src/thir.rs
+++ b/compiler/rustc_middle/src/thir.rs
@@ -20,7 +20,7 @@ use rustc_middle::mir::interpret::AllocId;
 use rustc_middle::mir::{self, BinOp, BorrowKind, FakeReadCause, Field, Mutability, UnOp};
 use rustc_middle::ty::adjustment::PointerCast;
 use rustc_middle::ty::subst::SubstsRef;
-use rustc_middle::ty::{self, AdtDef, Ty, UpvarSubsts};
+use rustc_middle::ty::{self, AdtDef, FnSig, Ty, UpvarSubsts};
 use rustc_middle::ty::{CanonicalUserType, CanonicalUserTypeAnnotation};
 use rustc_span::def_id::LocalDefId;
 use rustc_span::{sym, Span, Symbol, DUMMY_SP};
@@ -33,7 +33,12 @@ pub mod print;
 pub mod visit;
 
 macro_rules! thir_with_elements {
-    ($($name:ident: $id:ty => $value:ty => $format:literal,)*) => {
+    (
+        $($field_name:ident: $field_ty:ty,)*
+
+    @elements:
+        $($name:ident: $id:ty => $value:ty => $format:literal,)*
+    ) => {
         $(
             newtype_index! {
                 #[derive(HashStable)]
@@ -48,14 +53,20 @@ macro_rules! thir_with_elements {
         #[derive(Debug, HashStable, Clone)]
         pub struct Thir<'tcx> {
             $(
+                pub $field_name: $field_ty,
+            )*
+            $(
                 pub $name: IndexVec<$id, $value>,
             )*
         }
 
         impl<'tcx> Thir<'tcx> {
-            pub fn new() -> Thir<'tcx> {
+            pub fn new($($field_name: $field_ty,)*) -> Thir<'tcx> {
                 Thir {
                     $(
+                        $field_name,
+                    )*
+                    $(
                         $name: IndexVec::new(),
                     )*
                 }
@@ -76,6 +87,9 @@ macro_rules! thir_with_elements {
 pub const UPVAR_ENV_PARAM: ParamId = ParamId::from_u32(0);
 
 thir_with_elements! {
+    body_type: BodyTy<'tcx>,
+
+@elements:
     arms: ArmId => Arm<'tcx> => "a{}",
     blocks: BlockId => Block => "b{}",
     exprs: ExprId => Expr<'tcx> => "e{}",
@@ -83,6 +97,12 @@ thir_with_elements! {
     params: ParamId => Param<'tcx> => "p{}",
 }
 
+#[derive(Debug, HashStable, Clone)]
+pub enum BodyTy<'tcx> {
+    Const(Ty<'tcx>),
+    Fn(FnSig<'tcx>),
+}
+
 /// Description of a type-checked function parameter.
 #[derive(Clone, Debug, HashStable)]
 pub struct Param<'tcx> {
diff --git a/compiler/rustc_mir_build/src/build/mod.rs b/compiler/rustc_mir_build/src/build/mod.rs
index a6de8684c0f..9232c931afb 100644
--- a/compiler/rustc_mir_build/src/build/mod.rs
+++ b/compiler/rustc_mir_build/src/build/mod.rs
@@ -47,8 +47,6 @@ pub(crate) fn mir_built(
 
 /// Construct the MIR for a given `DefId`.
 fn mir_build(tcx: TyCtxt<'_>, def: ty::WithOptConstParam<LocalDefId>) -> Body<'_> {
-    let body_owner_kind = tcx.hir().body_owner_kind(def.did);
-
     // Ensure unsafeck and abstract const building is ran before we steal the THIR.
     // We can't use `ensure()` for `thir_abstract_const` as it doesn't compute the query
     // if inputs are green. This can cause ICEs when calling `thir_abstract_const` after
@@ -65,16 +63,15 @@ fn mir_build(tcx: TyCtxt<'_>, def: ty::WithOptConstParam<LocalDefId>) -> Body<'_
     }
 
     let body = match tcx.thir_body(def) {
-        Err(error_reported) => construct_error(tcx, def.did, body_owner_kind, error_reported),
+        Err(error_reported) => construct_error(tcx, def.did, error_reported),
         Ok((thir, expr)) => {
             // We ran all queries that depended on THIR at the beginning
             // of `mir_build`, so now we can steal it
             let thir = thir.steal();
 
-            if body_owner_kind.is_fn_or_closure() {
-                construct_fn(tcx, def, &thir, expr)
-            } else {
-                construct_const(tcx, def, &thir, expr)
+            match thir.body_type {
+                thir::BodyTy::Fn(fn_sig) => construct_fn(tcx, def, &thir, expr, fn_sig),
+                thir::BodyTy::Const(ty) => construct_const(tcx, def, &thir, expr, ty),
             }
         }
     };
@@ -434,6 +431,7 @@ fn construct_fn<'tcx>(
     fn_def: ty::WithOptConstParam<LocalDefId>,
     thir: &Thir<'tcx>,
     expr: ExprId,
+    fn_sig: ty::FnSig<'tcx>,
 ) -> Body<'tcx> {
     let span = tcx.def_span(fn_def.did);
     let fn_id = tcx.hir().local_def_id_to_hir_id(fn_def.did);
@@ -453,11 +451,6 @@ fn construct_fn<'tcx>(
         .output
         .span();
 
-    // fetch the fully liberated fn signature (that is, all bound
-    // types/lifetimes replaced)
-    let typeck_results = tcx.typeck_opt_const_arg(fn_def);
-    let fn_sig = typeck_results.liberated_fn_sigs()[fn_id];
-
     let safety = match fn_sig.unsafety {
         hir::Unsafety::Normal => Safety::Safe,
         hir::Unsafety::Unsafe => Safety::FnUnsafe,
@@ -563,6 +556,7 @@ fn construct_const<'a, 'tcx>(
     def: ty::WithOptConstParam<LocalDefId>,
     thir: &'a Thir<'tcx>,
     expr: ExprId,
+    const_ty: Ty<'tcx>,
 ) -> Body<'tcx> {
     let hir_id = tcx.hir().local_def_id_to_hir_id(def.did);
 
@@ -586,20 +580,6 @@ fn construct_const<'a, 'tcx>(
         _ => span_bug!(tcx.def_span(def.did), "can't build MIR for {:?}", def.did),
     };
 
-    // Get the revealed type of this const. This is *not* the adjusted
-    // type of its body, which may be a subtype of this type. For
-    // example:
-    //
-    // fn foo(_: &()) {}
-    // static X: fn(&'static ()) = foo;
-    //
-    // The adjusted type of the body of X is `for<'a> fn(&'a ())` which
-    // is not the same as the type of X. We need the type of the return
-    // place to be the type of the constant because NLL typeck will
-    // equate them.
-    let typeck_results = tcx.typeck_opt_const_arg(def);
-    let const_ty = typeck_results.node_type(hir_id);
-
     let infcx = tcx.infer_ctxt().build();
     let mut builder = Builder::new(
         thir,
@@ -629,15 +609,11 @@ fn construct_const<'a, 'tcx>(
 ///
 /// This is required because we may still want to run MIR passes on an item
 /// with type errors, but normal MIR construction can't handle that in general.
-fn construct_error(
-    tcx: TyCtxt<'_>,
-    def: LocalDefId,
-    body_owner_kind: hir::BodyOwnerKind,
-    err: ErrorGuaranteed,
-) -> Body<'_> {
+fn construct_error(tcx: TyCtxt<'_>, def: LocalDefId, err: ErrorGuaranteed) -> Body<'_> {
     let span = tcx.def_span(def);
     let hir_id = tcx.hir().local_def_id_to_hir_id(def);
     let generator_kind = tcx.generator_kind(def);
+    let body_owner_kind = tcx.hir().body_owner_kind(def);
 
     let ty = tcx.ty_error(err);
     let num_params = match body_owner_kind {
diff --git a/compiler/rustc_mir_build/src/thir/cx/mod.rs b/compiler/rustc_mir_build/src/thir/cx/mod.rs
index 20af60a511e..39a68048e26 100644
--- a/compiler/rustc_mir_build/src/thir/cx/mod.rs
+++ b/compiler/rustc_mir_build/src/thir/cx/mod.rs
@@ -99,9 +99,30 @@ impl<'tcx> Cx<'tcx> {
         let typeck_results = tcx.typeck_opt_const_arg(def);
         let did = def.did;
         let hir = tcx.hir();
+        let hir_id = hir.local_def_id_to_hir_id(did);
+
+        let body_type = if hir.body_owner_kind(did).is_fn_or_closure() {
+            // fetch the fully liberated fn signature (that is, all bound
+            // types/lifetimes replaced)
+            BodyTy::Fn(typeck_results.liberated_fn_sigs()[hir_id])
+        } else {
+            // Get the revealed type of this const. This is *not* the adjusted
+            // type of its body, which may be a subtype of this type. For
+            // example:
+            //
+            // fn foo(_: &()) {}
+            // static X: fn(&'static ()) = foo;
+            //
+            // The adjusted type of the body of X is `for<'a> fn(&'a ())` which
+            // is not the same as the type of X. We need the type of the return
+            // place to be the type of the constant because NLL typeck will
+            // equate them.
+            BodyTy::Const(typeck_results.node_type(hir_id))
+        };
+
         Cx {
             tcx,
-            thir: Thir::new(),
+            thir: Thir::new(body_type),
             param_env: tcx.param_env(def.did),
             region_scope_tree: tcx.region_scope_tree(def.did),
             typeck_results,
@@ -109,7 +130,7 @@ impl<'tcx> Cx<'tcx> {
             body_owner: did.to_def_id(),
             adjustment_span: None,
             apply_adjustments: hir
-                .attrs(hir.local_def_id_to_hir_id(did))
+                .attrs(hir_id)
                 .iter()
                 .all(|attr| attr.name_or_empty() != rustc_span::sym::custom_mir),
         }
diff --git a/tests/ui/thir-print/thir-flat.stdout b/tests/ui/thir-print/thir-flat.stdout
index c399fa66b6a..c31e6a218ce 100644
--- a/tests/ui/thir-print/thir-flat.stdout
+++ b/tests/ui/thir-print/thir-flat.stdout
@@ -1,5 +1,8 @@
 DefId(0:3 ~ thir_flat[45a6]::main):
 Thir {
+    body_type: Fn(
+        ([]; c_variadic: false)->(),
+    ),
     arms: [],
     blocks: [
         Block {