add Yield and Return type into generator

3 files changed, 56 insertions, 9 deletions

diff --git a/src/librustc/ty/sty.rs b/src/librustc/ty/sty.rs
index e581eac9ccd..436238b5e37 100644
--- a/src/librustc/ty/sty.rs
+++ b/src/librustc/ty/sty.rs
@@ -256,6 +256,17 @@ pub enum TypeVariants<'tcx> {
 /// closure C wind up influencing the decisions we ought to make for
 /// closure C (which would then require fixed point iteration to
 /// handle). Plus it fixes an ICE. :P
+///
+/// ## Generators
+///
+/// Perhaps surprisingly, `ClosureSubsts` are also used for
+/// generators.  In that case, what is written above is only half-true
+/// -- the set of type parameters is similar, but the role of CK and
+/// CS are different.  CK represents the "yield type" and CS
+/// represents the "return type" of the generator.
+///
+/// It'd be nice to split this struct into ClosureSubsts and
+/// GeneratorSubsts, I believe. -nmatsakis
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug, RustcEncodable, RustcDecodable)]
 pub struct ClosureSubsts<'tcx> {
     /// Lifetime and type parameters from the enclosing function,
@@ -310,6 +321,31 @@ impl<'tcx> ClosureSubsts<'tcx> {
     pub fn closure_sig_ty(self, def_id: DefId, tcx: TyCtxt<'_, '_, '_>) -> Ty<'tcx> {
         self.split(def_id, tcx).closure_sig_ty
     }
+
+    /// Returns the type representing the yield type of the generator.
+    pub fn generator_yield_ty(self, def_id: DefId, tcx: TyCtxt<'_, '_, '_>) -> Ty<'tcx> {
+        self.closure_kind_ty(def_id, tcx)
+    }
+
+    /// Returns the type representing the return type of the generator.
+    pub fn generator_return_ty(self, def_id: DefId, tcx: TyCtxt<'_, '_, '_>) -> Ty<'tcx> {
+        self.closure_sig_ty(def_id, tcx)
+    }
+
+    /// Return the "generator signature", which consists of its yield
+    /// and return types.
+    ///
+    /// NB. We treat this as a `PolyGenSig`, but since it only
+    /// contains associated types of the generator, at present it
+    /// never binds any regions.
+    pub fn generator_poly_sig(self, def_id: DefId, tcx: TyCtxt<'_, '_, '_>) -> PolyGenSig<'tcx> {
+        ty::Binder(
+            ty::GenSig {
+                yield_ty: self.generator_yield_ty(def_id, tcx),
+                return_ty: self.generator_return_ty(def_id, tcx),
+            }
+        )
+    }
 }
 
 impl<'tcx> ClosureSubsts<'tcx> {
diff --git a/src/librustc_typeck/check/closure.rs b/src/librustc_typeck/check/closure.rs
index 5fd1a0afb3e..5b5d697bcf4 100644
--- a/src/librustc_typeck/check/closure.rs
+++ b/src/librustc_typeck/check/closure.rs
@@ -10,7 +10,7 @@
 
 //! Code for type-checking closure expressions.
 
-use super::{check_fn, Expectation, FnCtxt};
+use super::{check_fn, Expectation, FnCtxt, GeneratorTypes};
 
 use astconv::AstConv;
 use rustc::hir::def_id::DefId;
@@ -79,7 +79,7 @@ impl<'a, 'gcx, 'tcx> FnCtxt<'a, 'gcx, 'tcx> {
 
         debug!("check_closure: ty_of_closure returns {:?}", liberated_sig);
 
-        let interior = check_fn(
+        let generator_types = check_fn(
             self,
             self.param_env,
             liberated_sig,
@@ -106,13 +106,13 @@ impl<'a, 'gcx, 'tcx> FnCtxt<'a, 'gcx, 'tcx> {
         let substs = ty::ClosureSubsts { substs };
         let closure_type = self.tcx.mk_closure(expr_def_id, substs);
 
-        if let Some(interior) = interior {
+        if let Some(GeneratorTypes { yield_ty, interior }) = generator_types {
             self.demand_eqtype(expr.span,
-                               ty::ClosureKind::FnOnce.to_ty(self.tcx),
-                               substs.closure_kind_ty(expr_def_id, self.tcx));
+                               yield_ty,
+                               substs.generator_yield_ty(expr_def_id, self.tcx));
             self.demand_eqtype(expr.span,
-                               self.tcx.types.char, // for generator, use some bogus type
-                               substs.closure_sig_ty(expr_def_id, self.tcx));
+                               liberated_sig.output(),
+                               substs.generator_return_ty(expr_def_id, self.tcx));
             return self.tcx.mk_generator(expr_def_id, substs, interior);
         }
 
diff --git a/src/librustc_typeck/check/mod.rs b/src/librustc_typeck/check/mod.rs
index ce8c285f0e2..91283ab7a2f 100644
--- a/src/librustc_typeck/check/mod.rs
+++ b/src/librustc_typeck/check/mod.rs
@@ -972,6 +972,17 @@ impl<'a, 'gcx, 'tcx> Visitor<'gcx> for GatherLocalsVisitor<'a, 'gcx, 'tcx> {
                 _: hir::BodyId, _: Span, _: ast::NodeId) { }
 }
 
+/// When `check_fn` is invoked on a generator (i.e., a body that
+/// includes yield), it returns back some information about the yield
+/// points.
+struct GeneratorTypes<'tcx> {
+    /// Type of value that is yielded.
+    yield_ty: ty::Ty<'tcx>,
+
+    /// Types that are captured (see `GeneratorInterior` for more).
+    interior: ty::GeneratorInterior<'tcx>
+}
+
 /// Helper used for fns and closures. Does the grungy work of checking a function
 /// body and returns the function context used for that purpose, since in the case of a fn item
 /// there is still a bit more to do.
@@ -985,7 +996,7 @@ fn check_fn<'a, 'gcx, 'tcx>(inherited: &'a Inherited<'a, 'gcx, 'tcx>,
                             fn_id: ast::NodeId,
                             body: &'gcx hir::Body,
                             can_be_generator: bool)
-                            -> (FnCtxt<'a, 'gcx, 'tcx>, Option<ty::GeneratorInterior<'tcx>>)
+                            -> (FnCtxt<'a, 'gcx, 'tcx>, Option<GeneratorTypes<'tcx>>)
 {
     let mut fn_sig = fn_sig.clone();
 
@@ -1047,7 +1058,7 @@ fn check_fn<'a, 'gcx, 'tcx>(inherited: &'a Inherited<'a, 'gcx, 'tcx>,
 
         inherited.tables.borrow_mut().generator_interiors_mut().insert(fn_hir_id, interior);
 
-        Some(interior)
+        Some(GeneratorTypes { yield_ty: gen_sig.yield_ty, interior: interior })
     } else {
         inherited.tables.borrow_mut().generator_sigs_mut().insert(fn_hir_id, None);
         None