Generalize discriminant info calls for generators and ADTs

8 files changed, 95 insertions, 51 deletions

diff --git a/src/librustc/ty/mod.rs b/src/librustc/ty/mod.rs
index 67be228d232..82696c22a22 100644
--- a/src/librustc/ty/mod.rs
+++ b/src/librustc/ty/mod.rs
@@ -41,6 +41,7 @@ use std::ops::Deref;
 use rustc_data_structures::sync::{self, Lrc, ParallelIterator, par_iter};
 use std::slice;
 use std::{mem, ptr};
+use std::ops::Range;
 use syntax::ast::{self, Name, Ident, NodeId};
 use syntax::attr;
 use syntax::ext::hygiene::Mark;
@@ -2418,11 +2419,17 @@ impl<'a, 'gcx, 'tcx> AdtDef {
         })
     }
 
+    #[inline]
+    pub fn variant_range(&self) -> Range<VariantIdx> {
+        (VariantIdx::new(0)..VariantIdx::new(self.variants.len()))
+    }
+
     /// Computes the discriminant value used by a specific variant.
     /// Unlike `discriminants`, this is (amortized) constant-time,
     /// only doing at most one query for evaluating an explicit
     /// discriminant (the last one before the requested variant),
     /// assuming there are no constant-evaluation errors there.
+    #[inline]
     pub fn discriminant_for_variant(&self,
                                     tcx: TyCtxt<'a, 'gcx, 'tcx>,
                                     variant_index: VariantIdx)
diff --git a/src/librustc/ty/sty.rs b/src/librustc/ty/sty.rs
index af447fc8a48..331c54d1d46 100644
--- a/src/librustc/ty/sty.rs
+++ b/src/librustc/ty/sty.rs
@@ -9,7 +9,7 @@ use polonius_engine::Atom;
 use rustc_data_structures::indexed_vec::Idx;
 use rustc_macros::HashStable;
 use crate::ty::subst::{InternalSubsts, Subst, SubstsRef, Kind, UnpackedKind};
-use crate::ty::{self, AdtDef, DefIdTree, TypeFlags, Ty, TyCtxt, TypeFoldable};
+use crate::ty::{self, AdtDef, Discr, DefIdTree, TypeFlags, Ty, TyCtxt, TypeFoldable};
 use crate::ty::{List, TyS, ParamEnvAnd, ParamEnv};
 use crate::ty::layout::VariantIdx;
 use crate::util::captures::Captures;
@@ -18,6 +18,7 @@ use crate::mir::interpret::{Scalar, Pointer};
 use smallvec::SmallVec;
 use std::cmp::Ordering;
 use std::marker::PhantomData;
+use std::ops::Range;
 use rustc_target::spec::abi;
 use syntax::ast::{self, Ident};
 use syntax::symbol::{keywords, InternedString};
@@ -478,14 +479,35 @@ impl<'a, 'gcx, 'tcx> GeneratorSubsts<'tcx> {
     const RETURNED_NAME: &'static str = "Returned";
     const POISONED_NAME: &'static str = "Panicked";
 
-    /// The variants of this Generator.
+    /// The valid variant indices of this Generator.
     #[inline]
-    pub fn variants(&self, def_id: DefId, tcx: TyCtxt<'a, 'gcx, 'tcx>) ->
-        impl Iterator<Item = VariantIdx>
-    {
+    pub fn variant_range(&self, def_id: DefId, tcx: TyCtxt<'a, 'gcx, 'tcx>) -> Range<VariantIdx> {
         // FIXME requires optimized MIR
         let num_variants = self.state_tys(def_id, tcx).count();
-        (0..num_variants).map(VariantIdx::new)
+        (VariantIdx::new(0)..VariantIdx::new(num_variants))
+    }
+
+    /// The discriminant for the given variant. Panics if the variant_index is
+    /// out of range.
+    #[inline]
+    pub fn discriminant_for_variant(
+        &self, def_id: DefId, tcx: TyCtxt<'a, 'gcx, 'tcx>, variant_index: VariantIdx
+    ) -> Discr<'tcx> {
+        // Generators don't support explicit discriminant values, so they are
+        // the same as the variant index.
+        assert!(self.variant_range(def_id, tcx).contains(&variant_index));
+        Discr { val: variant_index.as_usize() as u128, ty: self.discr_ty(tcx) }
+    }
+
+    /// The set of all discriminants for the Generator, enumerated with their
+    /// variant indices.
+    #[inline]
+    pub fn discriminants(
+        &'a self, def_id: DefId, tcx: TyCtxt<'a, 'gcx, 'tcx>
+    ) -> impl Iterator<Item=(VariantIdx, Discr<'tcx>)> + Captures<'gcx> + 'a {
+        self.variant_range(def_id, tcx).map(move |index| {
+            (index, Discr { val: index.as_usize() as u128, ty: self.discr_ty(tcx) })
+        })
     }
 
     /// Calls `f` with a reference to the name of the enumerator for the given
@@ -503,7 +525,7 @@ impl<'a, 'gcx, 'tcx> GeneratorSubsts<'tcx> {
         f(name)
     }
 
-    /// The type of the state "discriminant" used in the generator type.
+    /// The type of the state discriminant used in the generator type.
     #[inline]
     pub fn discr_ty(&self, tcx: TyCtxt<'a, 'gcx, 'tcx>) -> Ty<'tcx> {
         tcx.types.u32
@@ -2028,6 +2050,34 @@ impl<'a, 'gcx, 'tcx> TyS<'tcx> {
         }
     }
 
+    /// If the type contains variants, returns the valid range of variant indices.
+    /// FIXME This requires the optimized MIR in the case of generators.
+    #[inline]
+    pub fn variant_range(&self, tcx: TyCtxt<'a, 'gcx, 'tcx>) -> Option<Range<VariantIdx>> {
+        match self.sty {
+            TyKind::Adt(adt, _) => Some(adt.variant_range()),
+            TyKind::Generator(def_id, substs, _) => Some(substs.variant_range(def_id, tcx)),
+            _ => None,
+        }
+    }
+
+    /// If the type contains variants, returns the variant for `variant_index`.
+    /// Panics if `variant_index` is out of range.
+    /// FIXME This requires the optimized MIR in the case of generators.
+    #[inline]
+    pub fn discriminant_for_variant(
+        &self,
+        tcx: TyCtxt<'a, 'gcx, 'tcx>,
+        variant_index: VariantIdx
+    ) -> Option<Discr<'tcx>> {
+        match self.sty {
+            TyKind::Adt(adt, _) => Some(adt.discriminant_for_variant(tcx, variant_index)),
+            TyKind::Generator(def_id, substs, _) =>
+                Some(substs.discriminant_for_variant(def_id, tcx, variant_index)),
+            _ => None,
+        }
+    }
+
     /// Push onto `out` the regions directly referenced from this type (but not
     /// types reachable from this type via `walk_tys`). This ignores late-bound
     /// regions binders.
diff --git a/src/librustc_codegen_llvm/debuginfo/metadata.rs b/src/librustc_codegen_llvm/debuginfo/metadata.rs
index 54d15e527c9..6f427f29fcc 100644
--- a/src/librustc_codegen_llvm/debuginfo/metadata.rs
+++ b/src/librustc_codegen_llvm/debuginfo/metadata.rs
@@ -1382,12 +1382,6 @@ impl EnumMemberDescriptionFactory<'ll, 'tcx> {
                                                   variant_type_metadata,
                                                   member_descriptions);
 
-                    // TODO make this into a helper
-                    let discriminant = match &self.layout.ty.sty {
-                        ty::Adt(adt, _) => adt.discriminant_for_variant(cx.tcx, i).val as u64,
-                        ty::Generator(..) => i.as_usize() as u64,
-                        _ => bug!(),
-                    }.into();
                     MemberDescription {
                         name: if fallback {
                             String::new()
@@ -1399,7 +1393,9 @@ impl EnumMemberDescriptionFactory<'ll, 'tcx> {
                         size: self.layout.size,
                         align: self.layout.align.abi,
                         flags: DIFlags::FlagZero,
-                        discriminant,
+                        discriminant: Some(
+                            self.layout.ty.discriminant_for_variant(cx.tcx, i).unwrap().val as u64
+                        ),
                     }
                 }).collect()
             }
@@ -1722,7 +1718,7 @@ fn prepare_enum_metadata(
                 })
                 .collect(),
             ty::Generator(_, substs, _) => substs
-                .variants(enum_def_id, cx.tcx)
+                .variant_range(enum_def_id, cx.tcx)
                 .map(|v| substs.map_variant_name(v, |name| {
                     let name = SmallCStr::new(name);
                     unsafe {
diff --git a/src/librustc_codegen_ssa/mir/place.rs b/src/librustc_codegen_ssa/mir/place.rs
index 2875468127e..670b6c47269 100644
--- a/src/librustc_codegen_ssa/mir/place.rs
+++ b/src/librustc_codegen_ssa/mir/place.rs
@@ -218,9 +218,8 @@ impl<'a, 'tcx: 'a, V: CodegenObject> PlaceRef<'tcx, V> {
         }
         let (discr_scalar, discr_kind, discr_index) = match self.layout.variants {
             layout::Variants::Single { index } => {
-                let discr_val = self.layout.ty.ty_adt_def().map_or(
-                    index.as_u32() as u128,
-                    |def| def.discriminant_for_variant(bx.cx().tcx(), index).val);
+                let discr_val = self.layout.ty.discriminant_for_variant(bx.cx().tcx(), index)
+                    .map_or(index.as_u32() as u128, |discr| discr.val);
                 return bx.cx().const_uint_big(cast_to, discr_val);
             }
             layout::Variants::Multiple { ref discr, ref discr_kind, discr_index, .. } => {
@@ -296,15 +295,8 @@ impl<'a, 'tcx: 'a, V: CodegenObject> PlaceRef<'tcx, V> {
                 ..
             } => {
                 let ptr = self.project_field(bx, discr_index);
-                let to = match self.layout.ty.sty {
-                    ty::TyKind::Adt(adt_def, _) => adt_def
-                        .discriminant_for_variant(bx.tcx(), variant_index)
-                        .val,
-                    // Generators don't support explicit discriminant values, so
-                    // they are the same as the variant index.
-                    ty::TyKind::Generator(..) => variant_index.as_u32() as u128,
-                    _ => bug!(),
-                };
+                let to =
+                    self.layout.ty.discriminant_for_variant(bx.tcx(), variant_index).unwrap().val;
                 bx.store(
                     bx.cx().const_uint_big(bx.cx().backend_type(ptr.layout), to),
                     ptr.llval,
diff --git a/src/librustc_codegen_ssa/mir/rvalue.rs b/src/librustc_codegen_ssa/mir/rvalue.rs
index 35e9d918aa6..8573021bef3 100644
--- a/src/librustc_codegen_ssa/mir/rvalue.rs
+++ b/src/librustc_codegen_ssa/mir/rvalue.rs
@@ -271,13 +271,12 @@ impl<'a, 'tcx: 'a, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         let ll_t_in = bx.cx().immediate_backend_type(operand.layout);
                         match operand.layout.variants {
                             layout::Variants::Single { index } => {
-                                if let Some(def) = operand.layout.ty.ty_adt_def() {
-                                    let discr_val = def
-                                        .discriminant_for_variant(bx.cx().tcx(), index)
-                                        .val;
-                                    let discr = bx.cx().const_uint_big(ll_t_out, discr_val);
+                                if let Some(discr) =
+                                    operand.layout.ty.discriminant_for_variant(bx.tcx(), index)
+                                {
+                                    let discr_val = bx.cx().const_uint_big(ll_t_out, discr.val);
                                     return (bx, OperandRef {
-                                        val: OperandValue::Immediate(discr),
+                                        val: OperandValue::Immediate(discr_val),
                                         layout: cast,
                                     });
                                 }
diff --git a/src/librustc_mir/interpret/cast.rs b/src/librustc_mir/interpret/cast.rs
index 32f218d49ce..58e474658ea 100644
--- a/src/librustc_mir/interpret/cast.rs
+++ b/src/librustc_mir/interpret/cast.rs
@@ -54,14 +54,13 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> InterpretCx<'a, 'mir, 'tcx, M>
                 } else {
                     match src.layout.variants {
                         layout::Variants::Single { index } => {
-                            if let Some(def) = src.layout.ty.ty_adt_def() {
+                            if let Some(discr) =
+                                src.layout.ty.discriminant_for_variant(*self.tcx, index)
+                            {
                                 // Cast from a univariant enum
                                 assert!(src.layout.is_zst());
-                                let discr_val = def
-                                    .discriminant_for_variant(*self.tcx, index)
-                                    .val;
                                 return self.write_scalar(
-                                    Scalar::from_uint(discr_val, dest.layout.size),
+                                    Scalar::from_uint(discr.val, dest.layout.size),
                                     dest);
                             }
                         }
diff --git a/src/librustc_mir/interpret/operand.rs b/src/librustc_mir/interpret/operand.rs
index 55c1bfb17de..334b22112fb 100644
--- a/src/librustc_mir/interpret/operand.rs
+++ b/src/librustc_mir/interpret/operand.rs
@@ -566,9 +566,9 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> InterpretCx<'a, 'mir, 'tcx, M>
 
         let (discr_kind, discr_index) = match rval.layout.variants {
             layout::Variants::Single { index } => {
-                let discr_val = rval.layout.ty.ty_adt_def().map_or(
+                let discr_val = rval.layout.ty.discriminant_for_variant(*self.tcx, index).map_or(
                     index.as_u32() as u128,
-                    |def| def.discriminant_for_variant(*self.tcx, index).val);
+                    |discr| discr.val);
                 return Ok((discr_val, index));
             }
             layout::Variants::Multiple { ref discr_kind, discr_index, .. } =>
@@ -603,12 +603,15 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> InterpretCx<'a, 'mir, 'tcx, M>
                     bits_discr
                 };
                 // Make sure we catch invalid discriminants
-                let index = rval.layout.ty
-                    .ty_adt_def()
-                    .expect("tagged layout for non adt")
-                    .discriminants(self.tcx.tcx)
-                    .find(|(_, var)| var.val == real_discr)
-                    .ok_or_else(|| InterpError::InvalidDiscriminant(raw_discr.erase_tag()))?;
+                let index = match &rval.layout.ty.sty {
+                    ty::Adt(adt, _) => adt
+                        .discriminants(self.tcx.tcx)
+                        .find(|(_, var)| var.val == real_discr),
+                    ty::Generator(def_id, substs, _) => substs
+                        .discriminants(*def_id, self.tcx.tcx)
+                        .find(|(_, var)| var.val == real_discr),
+                    _ => bug!("tagged layout for non-adt non-generator"),
+                }.ok_or_else(|| InterpError::InvalidDiscriminant(raw_discr.erase_tag()))?;
                 (real_discr, index.0)
             },
             layout::DiscriminantKind::Niche {
diff --git a/src/librustc_mir/interpret/place.rs b/src/librustc_mir/interpret/place.rs
index 8239337796e..0cc48002816 100644
--- a/src/librustc_mir/interpret/place.rs
+++ b/src/librustc_mir/interpret/place.rs
@@ -984,11 +984,9 @@ where
                 discr_index,
                 ..
             } => {
-                let adt_def = dest.layout.ty.ty_adt_def().unwrap();
-                assert!(variant_index.as_usize() < adt_def.variants.len());
-                let discr_val = adt_def
-                    .discriminant_for_variant(*self.tcx, variant_index)
-                    .val;
+                assert!(dest.layout.ty.variant_range(*self.tcx).unwrap().contains(&variant_index));
+                let discr_val =
+                    dest.layout.ty.discriminant_for_variant(*self.tcx, variant_index).unwrap().val;
 
                 // raw discriminants for enums are isize or bigger during
                 // their computation, but the in-memory tag is the smallest possible