Infer type of `yield` to be resume type

3 files changed, 24 insertions, 10 deletions

diff --git a/src/librustc_typeck/check/closure.rs b/src/librustc_typeck/check/closure.rs
index 97067e0b055..fd6be852051 100644
--- a/src/librustc_typeck/check/closure.rs
+++ b/src/librustc_typeck/check/closure.rs
@@ -92,11 +92,12 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 .into(),
             GenericParamDefKind::Const => span_bug!(expr.span, "closure has const param"),
         });
-        if let Some(GeneratorTypes { yield_ty, interior, movability }) = generator_types {
+        if let Some(GeneratorTypes { resume_ty, yield_ty, interior, movability }) = generator_types
+        {
             let generator_substs = substs.as_generator();
             self.demand_eqtype(
                 expr.span,
-                self.tcx.mk_unit(),  // WIP
+                resume_ty,
                 generator_substs.resume_ty(expr_def_id, self.tcx),
             );
             self.demand_eqtype(
diff --git a/src/librustc_typeck/check/expr.rs b/src/librustc_typeck/check/expr.rs
index b4c2b85241f..9ce89bd6363 100644
--- a/src/librustc_typeck/check/expr.rs
+++ b/src/librustc_typeck/check/expr.rs
@@ -1796,9 +1796,11 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         expr: &'tcx hir::Expr<'tcx>,
         src: &'tcx hir::YieldSource,
     ) -> Ty<'tcx> {
-        match self.yield_ty {
-            Some(ty) => {
-                self.check_expr_coercable_to_type(&value, ty);
+        match self.resume_yield_tys {
+            Some((resume_ty, yield_ty)) => {
+                self.check_expr_coercable_to_type(&value, yield_ty);
+
+                resume_ty
             }
             // Given that this `yield` expression was generated as a result of lowering a `.await`,
             // we know that the yield type must be `()`; however, the context won't contain this
@@ -1806,6 +1808,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             // value's type against `()` (this check should always hold).
             None if src == &hir::YieldSource::Await => {
                 self.check_expr_coercable_to_type(&value, self.tcx.mk_unit());
+                self.tcx.mk_unit()
             }
             _ => {
                 struct_span_err!(
@@ -1815,9 +1818,9 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                     "yield expression outside of generator literal"
                 )
                 .emit();
+                self.tcx.mk_unit()
             }
         }
-        self.tcx.mk_unit()
     }
 }
 
diff --git a/src/librustc_typeck/check/mod.rs b/src/librustc_typeck/check/mod.rs
index 0a917a1853e..9612500e3b0 100644
--- a/src/librustc_typeck/check/mod.rs
+++ b/src/librustc_typeck/check/mod.rs
@@ -573,7 +573,7 @@ pub struct FnCtxt<'a, 'tcx> {
     /// First span of a return site that we find. Used in error messages.
     ret_coercion_span: RefCell<Option<Span>>,
 
-    yield_ty: Option<Ty<'tcx>>,
+    resume_yield_tys: Option<(Ty<'tcx>, Ty<'tcx>)>,
 
     ps: RefCell<UnsafetyState>,
 
@@ -1248,6 +1248,9 @@ impl<'a, 'tcx> Visitor<'tcx> for GatherLocalsVisitor<'a, 'tcx> {
 /// includes yield), it returns back some information about the yield
 /// points.
 struct GeneratorTypes<'tcx> {
+    /// Type of generator argument / values returned by `yield`.
+    resume_ty: Ty<'tcx>,
+
     /// Type of value that is yielded.
     yield_ty: Ty<'tcx>,
 
@@ -1308,7 +1311,11 @@ fn check_fn<'a, 'tcx>(
         let yield_ty = fcx
             .next_ty_var(TypeVariableOrigin { kind: TypeVariableOriginKind::TypeInference, span });
         fcx.require_type_is_sized(yield_ty, span, traits::SizedYieldType);
-        fcx.yield_ty = Some(yield_ty);
+
+        // Resume type defaults to `()` if the generator has no argument.
+        let resume_ty = fn_sig.inputs().get(0).map(|ty| *ty).unwrap_or_else(|| tcx.mk_unit());
+
+        fcx.resume_yield_tys = Some((resume_ty, yield_ty));
     }
 
     let outer_def_id = tcx.closure_base_def_id(hir.local_def_id(fn_id));
@@ -1361,8 +1368,11 @@ fn check_fn<'a, 'tcx>(
         let interior = fcx
             .next_ty_var(TypeVariableOrigin { kind: TypeVariableOriginKind::MiscVariable, span });
         fcx.deferred_generator_interiors.borrow_mut().push((body.id(), interior, gen_kind));
+
+        let (resume_ty, yield_ty) = fcx.resume_yield_tys.unwrap();
         Some(GeneratorTypes {
-            yield_ty: fcx.yield_ty.unwrap(),
+            resume_ty,
+            yield_ty,
             interior,
             movability: can_be_generator.unwrap(),
         })
@@ -2764,7 +2774,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             err_count_on_creation: inh.tcx.sess.err_count(),
             ret_coercion: None,
             ret_coercion_span: RefCell::new(None),
-            yield_ty: None,
+            resume_yield_tys: None,
             ps: RefCell::new(UnsafetyState::function(hir::Unsafety::Normal, hir::CRATE_HIR_ID)),
             diverges: Cell::new(Diverges::Maybe),
             has_errors: Cell::new(false),