include `host_effect_index` in `Generics`

5 files changed, 34 insertions, 21 deletions

diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs
index 5681d365aa1..ab68436c093 100644
--- a/compiler/rustc_ast_lowering/src/item.rs
+++ b/compiler/rustc_ast_lowering/src/item.rs
@@ -417,14 +417,11 @@ impl<'hir> LoweringContext<'_, 'hir> {
             }
             ItemKind::Trait(box Trait { is_auto, unsafety, generics, bounds, items }) => {
                 // FIXME(const_trait_impl, effects, fee1-dead) this should be simplified if possible
-                let constness = if let Some(attrs) = attrs {
-                    attrs
-                        .iter()
-                        .find(|x| x.has_name(sym::const_trait))
-                        .map_or(Const::No, |x| Const::Yes(x.span))
-                } else {
-                    Const::No
-                };
+                let constness = attrs
+                    .unwrap_or(&[])
+                    .iter()
+                    .find(|x| x.has_name(sym::const_trait))
+                    .map_or(Const::No, |x| Const::Yes(x.span));
                 let (generics, (unsafety, items, bounds)) = self.lower_generics(
                     generics,
                     constness,
diff --git a/compiler/rustc_hir_analysis/src/collect/generics_of.rs b/compiler/rustc_hir_analysis/src/collect/generics_of.rs
index db446715596..ccc9f808411 100644
--- a/compiler/rustc_hir_analysis/src/collect/generics_of.rs
+++ b/compiler/rustc_hir_analysis/src/collect/generics_of.rs
@@ -101,6 +101,7 @@ pub(super) fn generics_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::Generics {
                         param_def_id_to_index,
                         has_self: generics.has_self,
                         has_late_bound_regions: generics.has_late_bound_regions,
+                        host_effect_index: None,
                     };
                 } else {
                     // HACK(eddyb) this provides the correct generics when
@@ -226,10 +227,12 @@ pub(super) fn generics_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::Generics {
     let has_self = opt_self.is_some();
     let mut parent_has_self = false;
     let mut own_start = has_self as u32;
+    let mut host_effect_index = None;
     let parent_count = parent_def_id.map_or(0, |def_id| {
         let generics = tcx.generics_of(def_id);
         assert!(!has_self);
         parent_has_self = generics.has_self;
+        host_effect_index = generics.host_effect_index;
         own_start = generics.count() as u32;
         generics.parent_count + generics.params.len()
     });
@@ -251,11 +254,11 @@ pub(super) fn generics_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::Generics {
 
     // Now create the real type and const parameters.
     let type_start = own_start - has_self as u32 + params.len() as u32;
-    let mut i = 0;
+    let mut i: u32 = 0;
     let mut next_index = || {
         let prev = i;
         i += 1;
-        prev as u32 + type_start
+        prev + type_start
     };
 
     const TYPE_DEFAULT_NOT_ALLOWED: &'static str = "defaults for type parameters are only allowed in \
@@ -295,10 +298,12 @@ pub(super) fn generics_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::Generics {
             })
         }
         GenericParamKind::Const { default, .. } => {
-            // `rustc_host` effect params are allowed to have defaults.
+            let is_host_param = tcx.has_attr(param.def_id, sym::rustc_host);
+
             if !matches!(allow_defaults, Defaults::Allowed)
                 && default.is_some()
-                && !tcx.has_attr(param.def_id, sym::rustc_host)
+                // `rustc_host` effect params are allowed to have defaults.
+                && !is_host_param
             {
                 tcx.sess.span_err(
                     param.span,
@@ -307,8 +312,18 @@ pub(super) fn generics_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::Generics {
                 );
             }
 
+            let index = next_index();
+
+            if is_host_param {
+                if let Some(idx) = host_effect_index {
+                    bug!("parent also has host effect param? index: {idx}, def: {def_id:?}");
+                }
+
+                host_effect_index = Some(parent_count + index as usize);
+            }
+
             Some(ty::GenericParamDef {
-                index: next_index(),
+                index,
                 name: param.name.ident().name,
                 def_id: param.def_id.to_def_id(),
                 pure_wrt_drop: param.pure_wrt_drop,
@@ -360,6 +375,7 @@ pub(super) fn generics_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::Generics {
         param_def_id_to_index,
         has_self: has_self || parent_has_self,
         has_late_bound_regions: has_late_bound_regions(tcx, node),
+        host_effect_index,
     }
 }
 
diff --git a/compiler/rustc_hir_typeck/src/callee.rs b/compiler/rustc_hir_typeck/src/callee.rs
index 7d6f6ad5e05..d29ba810d64 100644
--- a/compiler/rustc_hir_typeck/src/callee.rs
+++ b/compiler/rustc_hir_typeck/src/callee.rs
@@ -777,16 +777,11 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             None => tcx.consts.true_,
         };
 
-        let identity_substs = ty::InternalSubsts::identity_for_item(tcx, callee_did);
+        let generics = tcx.generics_of(callee_did);
 
-        trace!(?effect, ?identity_substs, ?callee_substs);
+        trace!(?effect, ?generics, ?callee_substs);
 
-        // FIXME this should be made more efficient
-        let host_effect_param_index = identity_substs.iter().position(|x| {
-            matches!(x.unpack(), ty::GenericArgKind::Const(const_) if matches!(const_.kind(), ty::ConstKind::Param(param) if param.name == sym::host))
-        });
-
-        if let Some(idx) = host_effect_param_index {
+        if let Some(idx) = generics.host_effect_index {
             let param = callee_substs.const_at(idx);
             let cause = self.misc(span);
             match self.at(&cause, self.param_env).eq(infer::DefineOpaqueTypes::No, effect, param) {
diff --git a/compiler/rustc_middle/src/ty/generics.rs b/compiler/rustc_middle/src/ty/generics.rs
index 58fd6e1aa27..ea82e0070b1 100644
--- a/compiler/rustc_middle/src/ty/generics.rs
+++ b/compiler/rustc_middle/src/ty/generics.rs
@@ -133,6 +133,9 @@ pub struct Generics {
 
     pub has_self: bool,
     pub has_late_bound_regions: Option<Span>,
+
+    // The index of the host effect when substituted. (i.e. might be index to parent substs)
+    pub host_effect_index: Option<usize>,
 }
 
 impl<'tcx> Generics {
diff --git a/compiler/rustc_ty_utils/src/assoc.rs b/compiler/rustc_ty_utils/src/assoc.rs
index b59458bbf35..1adaf5f572e 100644
--- a/compiler/rustc_ty_utils/src/assoc.rs
+++ b/compiler/rustc_ty_utils/src/assoc.rs
@@ -337,6 +337,7 @@ fn associated_type_for_impl_trait_in_trait(
             param_def_id_to_index,
             has_self: opaque_ty_generics.has_self,
             has_late_bound_regions: opaque_ty_generics.has_late_bound_regions,
+            host_effect_index: parent_generics.host_effect_index,
         }
     });
 
@@ -415,6 +416,7 @@ fn associated_type_for_impl_trait_in_impl(
             param_def_id_to_index,
             has_self: false,
             has_late_bound_regions: trait_assoc_generics.has_late_bound_regions,
+            host_effect_index: parent_generics.host_effect_index,
         }
     });