Update generator transform and generated function signature

4 files changed, 52 insertions, 0 deletions

diff --git a/src/libcore/pin.rs b/src/libcore/pin.rs
index 7c09a36d898..56a32c928fb 100644
--- a/src/libcore/pin.rs
+++ b/src/libcore/pin.rs
@@ -117,6 +117,7 @@ use ops::{Deref, DerefMut, Receiver, CoerceUnsized, DispatchFromDyn};
 // implementations, are allowed because they all only use `&P`, so they cannot move
 // the value behind `pointer`.
 #[stable(feature = "pin", since = "1.33.0")]
+#[cfg_attr(not(stage0), lang = "pin")]
 #[fundamental]
 #[repr(transparent)]
 #[derive(Copy, Clone, Hash, Eq, Ord)]
diff --git a/src/librustc/middle/lang_items.rs b/src/librustc/middle/lang_items.rs
index e324bde1f17..87107f727a0 100644
--- a/src/librustc/middle/lang_items.rs
+++ b/src/librustc/middle/lang_items.rs
@@ -300,6 +300,7 @@ language_item_table! {
     GeneratorStateLangItem,      "generator_state",    gen_state,               Target::Enum;
     GeneratorTraitLangItem,      "generator",          gen_trait,               Target::Trait;
     UnpinTraitLangItem,          "unpin",              unpin_trait,             Target::Trait;
+    PinTypeLangItem,             "pin",                pin_type,                Target::Struct;
 
     EqTraitLangItem,             "eq",                 eq_trait,                Target::Trait;
     PartialOrdTraitLangItem,     "partial_ord",        partial_ord_trait,       Target::Trait;
diff --git a/src/librustc/ty/instance.rs b/src/librustc/ty/instance.rs
index 78b6ffaa54e..bc43fedef0f 100644
--- a/src/librustc/ty/instance.rs
+++ b/src/librustc/ty/instance.rs
@@ -76,6 +76,11 @@ impl<'a, 'tcx> Instance<'tcx> {
                 let env_region = ty::ReLateBound(ty::INNERMOST, ty::BrEnv);
                 let env_ty = tcx.mk_mut_ref(tcx.mk_region(env_region), ty);
 
+                let pin_did = tcx.lang_items().pin_type().unwrap();
+                let pin_adt_ref = tcx.adt_def(pin_did);
+                let pin_substs = tcx.intern_substs(&[env_ty.into()]);
+                let env_ty = tcx.mk_adt(pin_adt_ref, pin_substs);
+
                 sig.map_bound(|sig| {
                     let state_did = tcx.lang_items().gen_state().unwrap();
                     let state_adt_ref = tcx.adt_def(state_did);
diff --git a/src/librustc_mir/transform/generator.rs b/src/librustc_mir/transform/generator.rs
index ec0c118634d..f5cc6a43e28 100644
--- a/src/librustc_mir/transform/generator.rs
+++ b/src/librustc_mir/transform/generator.rs
@@ -113,6 +113,33 @@ impl<'tcx> MutVisitor<'tcx> for DerefArgVisitor {
     }
 }
 
+struct PinArgVisitor<'tcx> {
+    ref_gen_ty: Ty<'tcx>,
+}
+
+impl<'tcx> MutVisitor<'tcx> for PinArgVisitor<'tcx> {
+    fn visit_local(&mut self,
+                   local: &mut Local,
+                   _: PlaceContext<'tcx>,
+                   _: Location) {
+        assert_ne!(*local, self_arg());
+    }
+
+    fn visit_place(&mut self,
+                    place: &mut Place<'tcx>,
+                    context: PlaceContext<'tcx>,
+                    location: Location) {
+        if *place == Place::Local(self_arg()) {
+            *place = Place::Projection(Box::new(Projection {
+                base: place.clone(),
+                elem: ProjectionElem::Field(Field::new(0), self.ref_gen_ty),
+            }));
+        } else {
+            self.super_place(place, context, location);
+        }
+    }
+}
+
 fn self_arg() -> Local {
     Local::new(1)
 }
@@ -286,6 +313,23 @@ fn make_generator_state_argument_indirect<'a, 'tcx>(
     DerefArgVisitor.visit_mir(mir);
 }
 
+fn make_generator_state_argument_pinned<'a, 'tcx>(
+                tcx: TyCtxt<'a, 'tcx, 'tcx>,
+                mir: &mut Mir<'tcx>) {
+    let ref_gen_ty = mir.local_decls.raw[1].ty;
+
+    let pin_did = tcx.lang_items().pin_type().unwrap();
+    let pin_adt_ref = tcx.adt_def(pin_did);
+    let substs = tcx.intern_substs(&[ref_gen_ty.into()]);
+    let pin_ref_gen_ty = tcx.mk_adt(pin_adt_ref, substs);
+
+    // Replace the by ref generator argument
+    mir.local_decls.raw[1].ty = pin_ref_gen_ty;
+
+    // Add the Pin field access to accesses of the generator state
+    PinArgVisitor { ref_gen_ty }.visit_mir(mir);
+}
+
 fn replace_result_variable<'tcx>(
     ret_ty: Ty<'tcx>,
     mir: &mut Mir<'tcx>,
@@ -741,6 +785,7 @@ fn create_generator_resume_function<'a, 'tcx>(
     insert_switch(tcx, mir, cases, &transform, TerminatorKind::Unreachable);
 
     make_generator_state_argument_indirect(tcx, def_id, mir);
+    make_generator_state_argument_pinned(tcx, mir);
 
     no_landing_pads(tcx, mir);