improve DiscriminantKind handling

src/libcore/marker.rs

@@ -693,6 +693,7 @@ mod impls {
 pub trait DiscriminantKind {
     /// The type of the discriminant, which must satisfy the trait
     /// bounds required by `mem::Discriminant`.
+    #[cfg_attr(not(bootstrap), lang = "discriminant_type")]
     type Discriminant: Clone + Copy + Debug + Eq + PartialEq + Hash + Send + Sync + Unpin;
 }
 

src/librustc_hir/lang_items.rs

@@ -193,6 +193,9 @@ language_item_table! {
     CloneTraitLangItem,            sym::clone,              clone_trait,             Target::Trait;
     SyncTraitLangItem,             sym::sync,               sync_trait,              Target::Trait;
     DiscriminantKindTraitLangItem, sym::discriminant_kind,  discriminant_kind_trait, Target::Trait;
+    // The associated item of `trait DiscriminantKind`.
+    DiscriminantTypeLangItem,      sym::discriminant_type,  discriminant_type,       Target::AssocTy;
+
     FreezeTraitLangItem,           sym::freeze,             freeze_trait,            Target::Trait;
 
     DropTraitLangItem,             sym::drop,               drop_trait,              Target::Trait;

src/librustc_span/symbol.rs

@@ -327,6 +327,7 @@ symbols! {
         diagnostic,
         direct,
         discriminant_kind,
+        discriminant_type,
         discriminant_value,
         dispatch_from_dyn,
         div,

src/librustc_trait_selection/traits/project.rs

@@ -23,11 +23,12 @@ use crate::traits::error_reporting::InferCtxtExt;
 use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_errors::ErrorReported;
 use rustc_hir::def_id::DefId;
-use rustc_hir::lang_items::{FnOnceOutputLangItem, FnOnceTraitLangItem, GeneratorTraitLangItem};
+use rustc_hir::lang_items::{
+    DiscriminantTypeLangItem, FnOnceOutputLangItem, FnOnceTraitLangItem, GeneratorTraitLangItem,
+};
 use rustc_infer::infer::resolve::OpportunisticRegionResolver;
 use rustc_middle::ty::fold::{TypeFoldable, TypeFolder};
 use rustc_middle::ty::subst::Subst;
-use rustc_middle::ty::util::IntTypeExt;
 use rustc_middle::ty::{self, ToPolyTraitRef, ToPredicate, Ty, TyCtxt, WithConstness};
 use rustc_span::symbol::sym;
 use rustc_span::DUMMY_SP;
@@ -1324,22 +1325,11 @@ fn confirm_discriminant_kind_candidate<'cx, 'tcx>(
     let self_ty = selcx.infcx().shallow_resolve(obligation.predicate.self_ty());
     let substs = tcx.mk_substs([self_ty.into()].iter());
 
-    let assoc_items = tcx.associated_items(tcx.lang_items().discriminant_kind_trait().unwrap());
-    // FIXME: emit an error if the trait definition is wrong
-    let discriminant_def_id = assoc_items.in_definition_order().next().unwrap().def_id;
-
-    let discriminant_ty = match self_ty.kind {
-        // Use the discriminant type for enums.
-        ty::Adt(adt, _) if adt.is_enum() => adt.repr.discr_type().to_ty(tcx),
-        // Default to `i32` for generators.
-        ty::Generator(..) => tcx.types.i32,
-        // Use `u8` for all other types.
-        _ => tcx.types.u8,
-    };
+    let discriminant_def_id = tcx.require_lang_item(DiscriminantTypeLangItem, None);
 
     let predicate = ty::ProjectionPredicate {
         projection_ty: ty::ProjectionTy { substs, item_def_id: discriminant_def_id },
-        ty: discriminant_ty,
+        ty: self_ty.discriminant_ty(tcx),
     };
 
     confirm_param_env_candidate(selcx, obligation, ty::Binder::bind(predicate))

src/test/ui/generator/discriminant.rs

@@ -66,8 +66,8 @@ macro_rules! yield250 {
 }
 
 fn cycle(
-    gen: impl Generator<()> + Unpin + DiscriminantKind<Discriminant = i32>,
-    expected_max_discr: i32
+    gen: impl Generator<()> + Unpin + DiscriminantKind<Discriminant = u32>,
+    expected_max_discr: u32
 ) {
     let mut gen = Box::pin(gen);
     let mut max_discr = 0;