Integrate generators to universal region setup

4 files changed, 29 insertions, 0 deletions

diff --git a/src/librustc/mir/visit.rs b/src/librustc/mir/visit.rs
index 58067931d56..57ed41f2f06 100644
--- a/src/librustc/mir/visit.rs
+++ b/src/librustc/mir/visit.rs
@@ -277,6 +277,13 @@ macro_rules! make_mir_visitor {
 
             fn super_mir(&mut self,
                          mir: & $($mutability)* Mir<'tcx>) {
+                if let Some(yield_ty) = &$($mutability)* mir.yield_ty {
+                    self.visit_ty(yield_ty, TyContext::YieldTy(SourceInfo {
+                        span: mir.span,
+                        scope: ARGUMENT_VISIBILITY_SCOPE,
+                    }));
+                }
+
                 // for best performance, we want to use an iterator rather
                 // than a for-loop, to avoid calling Mir::invalidate for
                 // each basic block.
@@ -852,6 +859,8 @@ pub enum TyContext {
     /// The return type of the function.
     ReturnTy(SourceInfo),
 
+    YieldTy(SourceInfo),
+
     /// A type found at some location.
     Location(Location),
 }
diff --git a/src/librustc_mir/borrow_check/nll/constraint_generation.rs b/src/librustc_mir/borrow_check/nll/constraint_generation.rs
index 7ca4ebd1cb2..3a39eb5c908 100644
--- a/src/librustc_mir/borrow_check/nll/constraint_generation.rs
+++ b/src/librustc_mir/borrow_check/nll/constraint_generation.rs
@@ -69,6 +69,7 @@ impl<'cg, 'cx, 'gcx, 'tcx> Visitor<'tcx> for ConstraintGeneration<'cg, 'cx, 'gcx
     fn visit_ty(&mut self, ty: &ty::Ty<'tcx>, ty_context: TyContext) {
         match ty_context {
             TyContext::ReturnTy(source_info) |
+            TyContext::YieldTy(source_info) |
             TyContext::LocalDecl { source_info, .. } => {
                 span_bug!(source_info.span,
                           "should not be visiting outside of the CFG: {:?}",
diff --git a/src/librustc_mir/borrow_check/nll/type_check/input_output.rs b/src/librustc_mir/borrow_check/nll/type_check/input_output.rs
index 9e88a632f5c..b1aeae0b76b 100644
--- a/src/librustc_mir/borrow_check/nll/type_check/input_output.rs
+++ b/src/librustc_mir/borrow_check/nll/type_check/input_output.rs
@@ -60,6 +60,15 @@ impl<'a, 'gcx, 'tcx> TypeChecker<'a, 'gcx, 'tcx> {
             self.equate_normalized_input_or_output(start_position, input_ty, mir_input_ty);
         }
 
+        assert!(
+            mir.yield_ty.is_some() && universal_regions.yield_ty.is_some() ||
+            mir.yield_ty.is_none() && universal_regions.yield_ty.is_none()
+            );
+        if let Some(mir_yield_ty) = mir.yield_ty {
+            let ur_yield_ty = universal_regions.yield_ty.unwrap();
+            self.equate_normalized_input_or_output(start_position, ur_yield_ty, mir_yield_ty);
+        }
+
         // Return types are a bit more complex. They may contain existential `impl Trait`
         // types.
         debug!(
diff --git a/src/librustc_mir/borrow_check/nll/universal_regions.rs b/src/librustc_mir/borrow_check/nll/universal_regions.rs
index 0e329266ea1..e426457a349 100644
--- a/src/librustc_mir/borrow_check/nll/universal_regions.rs
+++ b/src/librustc_mir/borrow_check/nll/universal_regions.rs
@@ -96,6 +96,8 @@ pub struct UniversalRegions<'tcx> {
     /// our special inference variable there, we would mess that up.
     pub region_bound_pairs: Vec<(ty::Region<'tcx>, GenericKind<'tcx>)>,
 
+    pub yield_ty: Option<Ty<'tcx>>,
+
     relations: UniversalRegionRelations,
 }
 
@@ -505,6 +507,13 @@ impl<'cx, 'gcx, 'tcx> UniversalRegionsBuilder<'cx, 'gcx, 'tcx> {
             num_universals
         );
 
+        let yield_ty = match defining_ty {
+            DefiningTy::Generator(def_id, substs, _) => {
+                Some(substs.generator_yield_ty(def_id, self.infcx.tcx))
+            }
+            _ => None,
+        };
+
         UniversalRegions {
             indices,
             fr_static,
@@ -516,6 +525,7 @@ impl<'cx, 'gcx, 'tcx> UniversalRegionsBuilder<'cx, 'gcx, 'tcx> {
             unnormalized_output_ty,
             unnormalized_input_tys,
             region_bound_pairs: self.region_bound_pairs,
+            yield_ty: yield_ty,
             relations: self.relations,
         }
     }