Change MIR building to fill in the resume place

This changes `Yield` from `as_rvalue` to `into` lowering, which could
have a possible performance impact. I could imagine special-casing
some resume types here to use a simpler lowering for them, but it's
unclear if that makes sense at this stage.

Author: jonas-schievink

src/librustc_mir_build/build/expr/as_rvalue.rs

@@ -230,18 +230,8 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 block = unpack!(this.stmt_expr(block, expr, None));
                 block.and(this.unit_rvalue())
             }
-            ExprKind::Yield { value } => {
-                let value = unpack!(block = this.as_operand(block, scope, value));
-                let resume = this.cfg.start_new_block();
-                let cleanup = this.generator_drop_cleanup();
-                this.cfg.terminate(
-                    block,
-                    source_info,
-                    TerminatorKind::Yield { value: value, resume: resume, drop: cleanup },
-                );
-                resume.and(this.unit_rvalue())
-            }
-            ExprKind::Literal { .. }
+            ExprKind::Yield { .. }
+            | ExprKind::Literal { .. }
             | ExprKind::StaticRef { .. }
             | ExprKind::Block { .. }
             | ExprKind::Match { .. }

src/librustc_mir_build/build/expr/category.rs

@@ -50,6 +50,7 @@ impl Category {
             | ExprKind::Adt { .. }
             | ExprKind::Borrow { .. }
             | ExprKind::AddressOf { .. }
+            | ExprKind::Yield { .. }
             | ExprKind::Call { .. } => Some(Category::Rvalue(RvalueFunc::Into)),
 
             ExprKind::Array { .. }
@@ -63,7 +64,6 @@ impl Category {
             | ExprKind::Repeat { .. }
             | ExprKind::Assign { .. }
             | ExprKind::AssignOp { .. }
-            | ExprKind::Yield { .. }
             | ExprKind::InlineAsm { .. } => Some(Category::Rvalue(RvalueFunc::AsRvalue)),
 
             ExprKind::Literal { .. } | ExprKind::StaticRef { .. } => Some(Category::Constant),

src/librustc_mir_build/build/expr/into.rs

@@ -365,6 +365,24 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 block.unit()
             }
 
+            ExprKind::Yield { value } => {
+                let scope = this.local_scope();
+                let value = unpack!(block = this.as_operand(block, scope, value));
+                let resume = this.cfg.start_new_block();
+                let cleanup = this.generator_drop_cleanup();
+                this.cfg.terminate(
+                    block,
+                    source_info,
+                    TerminatorKind::Yield {
+                        value,
+                        resume,
+                        resume_arg: destination.clone(),
+                        drop: cleanup,
+                    },
+                );
+                resume.unit()
+            }
+
             // these are the cases that are more naturally handled by some other mode
             ExprKind::Unary { .. }
             | ExprKind::Binary { .. }
@@ -376,8 +394,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             | ExprKind::Tuple { .. }
             | ExprKind::Closure { .. }
             | ExprKind::Literal { .. }
-            | ExprKind::StaticRef { .. }
-            | ExprKind::Yield { .. } => {
+            | ExprKind::StaticRef { .. } => {
                 debug_assert!(match Category::of(&expr.kind).unwrap() {
                     // should be handled above
                     Category::Rvalue(RvalueFunc::Into) => false,

src/librustc_mir_build/build/mod.rs

@@ -68,6 +68,12 @@ fn mir_build(tcx: TyCtxt<'_>, def_id: DefId) -> BodyAndCache<'_> {
             let fn_sig = cx.tables().liberated_fn_sigs()[id];
             let fn_def_id = tcx.hir().local_def_id(id);
 
+            let safety = match fn_sig.unsafety {
+                hir::Unsafety::Normal => Safety::Safe,
+                hir::Unsafety::Unsafe => Safety::FnUnsafe,
+            };
+
+            let body = tcx.hir().body(body_id);
             let ty = tcx.type_of(fn_def_id);
             let mut abi = fn_sig.abi;
             let implicit_argument = match ty.kind {
@@ -77,22 +83,23 @@ fn mir_build(tcx: TyCtxt<'_>, def_id: DefId) -> BodyAndCache<'_> {
                     abi = Abi::Rust;
                     vec![ArgInfo(liberated_closure_env_ty(tcx, id, body_id), None, None, None)]
                 }
-                ty::Generator(..) => {
+                ty::Generator(def_id, substs, _) => {
                     let gen_ty = tcx.body_tables(body_id).node_type(id);
-                    vec![
-                        ArgInfo(gen_ty, None, None, None),
-                        ArgInfo(tcx.mk_unit(), None, None, None),
-                    ]
+                    let resume_ty = substs.as_generator().resume_ty(def_id, tcx);
+
+                    // The resume argument may be missing, in that case we need to provide it here.
+                    if body.params.is_empty() {
+                        vec![
+                            ArgInfo(gen_ty, None, None, None),
+                            ArgInfo(resume_ty, None, None, None),
+                        ]
+                    } else {
+                        vec![ArgInfo(gen_ty, None, None, None)]
+                    }
                 }
                 _ => vec![],
             };
 
-            let safety = match fn_sig.unsafety {
-                hir::Unsafety::Normal => Safety::Safe,
-                hir::Unsafety::Unsafe => Safety::FnUnsafe,
-            };
-
-            let body = tcx.hir().body(body_id);
             let explicit_arguments = body.params.iter().enumerate().map(|(index, arg)| {
                 let owner_id = tcx.hir().body_owner(body_id);
                 let opt_ty_info;