Define generator discriminant type in only one place

Author: tmandry

src/librustc/ty/sty.rs

@@ -467,6 +467,11 @@ impl<'tcx> GeneratorSubsts<'tcx> {
 }
 
 impl<'a, 'gcx, 'tcx> GeneratorSubsts<'tcx> {
+    /// The type of the state "discriminant" used in the generator type.
+    pub fn discr_ty(&self, tcx: TyCtxt<'a, 'gcx, 'tcx>) -> Ty<'tcx> {
+        tcx.types.u32
+    }
+
     /// This returns the types of the MIR locals which had to be stored across suspension points.
     /// It is calculated in rustc_mir::transform::generator::StateTransform.
     /// All the types here must be in the tuple in GeneratorInterior.
@@ -481,15 +486,15 @@ impl<'a, 'gcx, 'tcx> GeneratorSubsts<'tcx> {
 
     /// This is the types of the fields of a generate which
     /// is available before the generator transformation.
-    /// It includes the upvars and the state discriminant which is u32.
+    /// It includes the upvars and the state discriminant.
     pub fn pre_transforms_tys(self, def_id: DefId, tcx: TyCtxt<'a, 'gcx, 'tcx>) ->
         impl Iterator<Item=Ty<'tcx>> + 'a
     {
-        self.upvar_tys(def_id, tcx).chain(iter::once(tcx.types.u32))
+        self.upvar_tys(def_id, tcx).chain(iter::once(self.discr_ty(tcx)))
     }
 
     /// This is the types of all the fields stored in a generator.
-    /// It includes the upvars, state types and the state discriminant which is u32.
+    /// It includes the upvars, state types and the state discriminant.
     pub fn field_tys(self, def_id: DefId, tcx: TyCtxt<'a, 'gcx, 'tcx>) ->
         impl Iterator<Item=Ty<'tcx>> + Captures<'gcx> + 'a
     {

src/librustc_mir/build/expr/as_rvalue.rs

@@ -255,15 +255,16 @@ impl<'a, 'gcx, 'tcx> Builder<'a, 'gcx, 'tcx> {
                         let movability = movability.unwrap();
                         // Add the state operand since it follows the upvars in the generator
                         // struct. See librustc_mir/transform/generator.rs for more details.
+                        let discr_ty = substs.discr_ty(this.hir.tcx());
                         operands.push(Operand::Constant(box Constant {
                             span: expr_span,
-                            ty: this.hir.tcx().types.u32,
+                            ty: discr_ty,
                             user_ty: None,
                             literal: this.hir.tcx().mk_const(
                                 ty::Const::from_bits(
                                     this.hir.tcx(),
                                     0,
-                                    ty::ParamEnv::empty().and(this.hir.tcx().types.u32),
+                                    ty::ParamEnv::empty().and(discr_ty),
                                 ),
                             ),
                         }));

src/librustc_mir/transform/generator.rs

@@ -166,6 +166,9 @@ struct TransformVisitor<'a, 'tcx: 'a> {
     // The index of the generator state in the generator struct
     state_field: usize,
 
+    // The type of the generator state in the generator struct
+    discr_ty: Ty<'tcx>,
+
     // Mapping from Local to (type of local, generator struct index)
     // FIXME(eddyb) This should use `IndexVec<Local, Option<_>>`.
     remap: FxHashMap<Local, (Ty<'tcx>, usize)>,
@@ -200,15 +203,15 @@ impl<'a, 'tcx> TransformVisitor<'a, 'tcx> {
 
     // Create a statement which changes the generator state
     fn set_state(&self, state_disc: u32, source_info: SourceInfo) -> Statement<'tcx> {
-        let state = self.make_field(self.state_field, self.tcx.types.u32);
+        let state = self.make_field(self.state_field, self.discr_ty);
         let val = Operand::Constant(box Constant {
             span: source_info.span,
-            ty: self.tcx.types.u32,
+            ty: self.discr_ty,
             user_ty: None,
             literal: self.tcx.mk_const(ty::Const::from_bits(
                 self.tcx,
                 state_disc.into(),
-                ty::ParamEnv::empty().and(self.tcx.types.u32)
+                ty::ParamEnv::empty().and(self.discr_ty)
             )),
         });
         Statement {
@@ -889,10 +892,11 @@ impl MirPass for StateTransform {
         let gen_ty = mir.local_decls.raw[1].ty;
 
         // Get the interior types and substs which typeck computed
-        let (upvars, interior, movable) = match gen_ty.sty {
+        let (upvars, interior, discr_ty, movable) = match gen_ty.sty {
             ty::Generator(_, substs, movability) => {
                 (substs.upvar_tys(def_id, tcx).collect(),
                  substs.witness(def_id, tcx),
+                 substs.discr_ty(tcx),
                  movability == hir::GeneratorMovability::Movable)
             }
             _ => bug!(),
@@ -937,6 +941,7 @@ impl MirPass for StateTransform {
             suspension_points: Vec::new(),
             new_ret_local,
             state_field,
+            discr_ty,
         };
         transform.visit_mir(mir);