Handle closures correctly in MIR inlining

Author: sinkuu

src/librustc_mir/transform/inline.rs

@@ -20,6 +20,7 @@ use rustc::mir::transform::{MirPass, MirSource};
 use rustc::mir::visit::*;
 use rustc::ty::{self, Instance, Ty, TyCtxt, TypeFoldable};
 use rustc::ty::subst::{Subst,Substs};
+use rustc::hir::map::definitions::DefPathData;
 
 use std::collections::VecDeque;
 use super::simplify::{remove_dead_blocks, CfgSimplifier};
@@ -550,22 +551,31 @@ impl<'a, 'tcx> Inliner<'a, 'tcx> {
         Operand::Consume(cast_tmp)
     }
 
-    fn make_call_args(&self, args: Vec<Operand<'tcx>>,
-                      callsite: &CallSite<'tcx>, caller_mir: &mut Mir<'tcx>) -> Vec<Operand<'tcx>> {
-        let tcx = self.tcx;
+    fn make_call_args(
+        &self,
+        args: Vec<Operand<'tcx>>,
+        callsite: &CallSite<'tcx>,
+        caller_mir: &mut Mir<'tcx>,
+    ) -> Vec<Operand<'tcx>> {
         // FIXME: Analysis of the usage of the arguments to avoid
         // unnecessary temporaries.
-        args.into_iter().map(|a| {
-            if let Operand::Consume(Lvalue::Local(local)) = a {
+
+        fn create_temp_if_necessary<'a, 'tcx: 'a>(
+            arg: Operand<'tcx>,
+            tcx: TyCtxt<'a, 'tcx, 'tcx>,
+            callsite: &CallSite<'tcx>,
+            caller_mir: &mut Mir<'tcx>,
+        ) -> Operand<'tcx> {
+            if let Operand::Consume(Lvalue::Local(local)) = arg {
                 if caller_mir.local_kind(local) == LocalKind::Temp {
                     // Reuse the operand if it's a temporary already
-                    return a;
+                    return arg;
                 }
             }
 
-            debug!("Creating temp for argument");
+            debug!("Creating temp for argument {:?}", arg);
             // Otherwise, create a temporary for the arg
-            let arg = Rvalue::Use(a);
+            let arg = Rvalue::Use(arg);
 
             let ty = arg.ty(caller_mir, tcx);
 
@@ -575,11 +585,47 @@ impl<'a, 'tcx> Inliner<'a, 'tcx> {
 
             let stmt = Statement {
                 source_info: callsite.location,
-                kind: StatementKind::Assign(arg_tmp.clone(), arg)
+                kind: StatementKind::Assign(arg_tmp.clone(), arg),
             };
             caller_mir[callsite.bb].statements.push(stmt);
             Operand::Consume(arg_tmp)
-        }).collect()
+        }
+
+        let tcx = self.tcx;
+
+        // A closure is passed its self-type and a tuple like `(arg1, arg2, ...)`,
+        // hence mappings to tuple fields are needed.
+        if tcx.def_key(callsite.callee).disambiguated_data.data == DefPathData::ClosureExpr {
+            let mut args = args.into_iter();
+
+            let self_ = create_temp_if_necessary(args.next().unwrap(), tcx, callsite, caller_mir);
+
+            let tuple = if let Operand::Consume(lvalue) =
+                create_temp_if_necessary(args.next().unwrap(), tcx, callsite, caller_mir)
+            {
+                lvalue
+            } else {
+                unreachable!()
+            };
+            assert!(args.next().is_none());
+
+            let tuple_tys = if let ty::TyTuple(s, _) = tuple.ty(caller_mir, tcx).to_ty(tcx).sty {
+                s
+            } else {
+                bug!("Closure arguments are not passed as a tuple");
+            };
+
+            let mut res = Vec::with_capacity(1 + tuple_tys.len());
+            res.push(self_);
+            res.extend(tuple_tys.iter().enumerate().map(|(i, ty)| {
+                Operand::Consume(tuple.clone().field(Field::new(i), ty))
+            }));
+            res
+        } else {
+            args.into_iter()
+                .map(|a| create_temp_if_necessary(a, tcx, callsite, caller_mir))
+                .collect()
+        }
     }
 }
 

src/test/mir-opt/inline-closure.rs

@@ -0,0 +1,53 @@
+// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// compile-flags: -Z span_free_formats
+
+// Tests that MIR inliner can handle closure arguments. (#45894)
+
+fn main() {
+    println!("{}", foo(0, 14));
+}
+
+fn foo<T: Copy>(_t: T, q: i32) -> i32 {
+    let x = |_t, _q| _t;
+    x(q*2, q*3)
+}
+
+// END RUST SOURCE
+// START rustc.foo.Inline.after.mir
+// ...
+// bb0: {
+//     StorageLive(_3);
+//     _3 = [closure@NodeId(28)];
+//     StorageLive(_4);
+//     _4 = &_3;
+//     StorageLive(_5);
+//     StorageLive(_6);
+//     StorageLive(_7);
+//     _7 = _2;
+//     _6 = Mul(_7, const 2i32);
+//     StorageDead(_7);
+//     StorageLive(_8);
+//     StorageLive(_9);
+//     _9 = _2;
+//     _8 = Mul(_9, const 3i32);
+//     StorageDead(_9);
+//     _5 = (_6, _8);
+//     _0 = (_5.0: i32);
+//     StorageDead(_5);
+//     StorageDead(_8);
+//     StorageDead(_6);
+//     StorageDead(_4);
+//     StorageDead(_3);
+//     return;
+// }
+// ...
+// END rustc.foo.Inline.after.mir