From fcf52c167f9280efe332557195998650a99bf437 Mon Sep 17 00:00:00 2001
From: Bastian Kauschke <bastian_kauschke@hotmail.de>
Date: Wed, 15 Jul 2020 09:59:08 +0200
Subject: [PATCH] improve DiscriminantKind handling

This now reuses `fn discriminant_ty` in project, removing
some code duplication. Doing so made me realize that
we previously had a disagreement about the discriminant
type of generators, with MIR using `u32` and codegen and
trait selection using `i32`.

We now always use `u32`.
---
 src/libcore/marker.rs                         |  1 +
 src/librustc_hir/lang_items.rs                |  3 +++
 src/librustc_span/symbol.rs                   |  1 +
 .../traits/project.rs                         | 20 +++++--------------
 src/test/ui/generator/discriminant.rs         |  4 ++--
 5 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/src/libcore/marker.rs b/src/libcore/marker.rs
index 6ddd41a26f122..56dddee7b7799 100644
--- a/src/libcore/marker.rs
+++ b/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;
 }
 
diff --git a/src/librustc_hir/lang_items.rs b/src/librustc_hir/lang_items.rs
index 1c8a56e5d80e7..88c97d874bed6 100644
--- a/src/librustc_hir/lang_items.rs
+++ b/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;
diff --git a/src/librustc_span/symbol.rs b/src/librustc_span/symbol.rs
index 0e9df5feb32ba..227cbbb4ca6ac 100644
--- a/src/librustc_span/symbol.rs
+++ b/src/librustc_span/symbol.rs
@@ -327,6 +327,7 @@ symbols! {
         diagnostic,
         direct,
         discriminant_kind,
+        discriminant_type,
         discriminant_value,
         dispatch_from_dyn,
         div,
diff --git a/src/librustc_trait_selection/traits/project.rs b/src/librustc_trait_selection/traits/project.rs
index bd86109e5a491..c08198ec373b8 100644
--- a/src/librustc_trait_selection/traits/project.rs
+++ b/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))
diff --git a/src/test/ui/generator/discriminant.rs b/src/test/ui/generator/discriminant.rs
index 3d0930da42243..195e77022992d 100644
--- a/src/test/ui/generator/discriminant.rs
+++ b/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;