Auto merge of #45913 - sinkuu:mir-inlining-closure, r=arielb1

Handle closures correctly in MIR inlining

Fixes #45894.

4 files changed, 113 insertions, 29 deletions

diff --git a/src/librustc/ty/util.rs b/src/librustc/ty/util.rs
index aa07a6b070e..a0219f2f95b 100644
--- a/src/librustc/ty/util.rs
+++ b/src/librustc/ty/util.rs
@@ -621,9 +621,13 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
         result
     }
 
+    pub fn is_closure(self, def_id: DefId) -> bool {
+        self.def_key(def_id).disambiguated_data.data == DefPathData::ClosureExpr
+    }
+
     pub fn closure_base_def_id(self, def_id: DefId) -> DefId {
         let mut def_id = def_id;
-        while self.def_key(def_id).disambiguated_data.data == DefPathData::ClosureExpr {
+        while self.is_closure(def_id) {
             def_id = self.parent_def_id(def_id).unwrap_or_else(|| {
                 bug!("closure {:?} has no parent", def_id);
             });
diff --git a/src/librustc_mir/transform/check_unsafety.rs b/src/librustc_mir/transform/check_unsafety.rs
index cbf1b0b0899..302e65a0e34 100644
--- a/src/librustc_mir/transform/check_unsafety.rs
+++ b/src/librustc_mir/transform/check_unsafety.rs
@@ -15,7 +15,6 @@ use rustc::ty::maps::Providers;
 use rustc::ty::{self, TyCtxt};
 use rustc::hir;
 use rustc::hir::def_id::DefId;
-use rustc::hir::map::DefPathData;
 use rustc::lint::builtin::{SAFE_EXTERN_STATICS, UNUSED_UNSAFE};
 use rustc::mir::*;
 use rustc::mir::visit::{LvalueContext, Visitor};
@@ -362,11 +361,11 @@ fn report_unused_unsafe(tcx: TyCtxt, used_unsafe: &FxHashSet<ast::NodeId>, id: a
 
 pub fn check_unsafety<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, def_id: DefId) {
     debug!("check_unsafety({:?})", def_id);
-    match tcx.def_key(def_id).disambiguated_data.data {
-        // closures are handled by their parent fn.
-        DefPathData::ClosureExpr => return,
-        _ => {}
-    };
+
+    // closures are handled by their parent fn.
+    if tcx.is_closure(def_id) {
+        return;
+    }
 
     let UnsafetyCheckResult {
         violations,
diff --git a/src/librustc_mir/transform/inline.rs b/src/librustc_mir/transform/inline.rs
index 5dec6177b36..628a8161615 100644
--- a/src/librustc_mir/transform/inline.rs
+++ b/src/librustc_mir/transform/inline.rs
@@ -550,36 +550,75 @@ 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>> {
+    fn make_call_args(
+        &self,
+        args: Vec<Operand<'tcx>>,
+        callsite: &CallSite<'tcx>,
+        caller_mir: &mut Mir<'tcx>,
+    ) -> Vec<Operand<'tcx>> {
         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.is_closure(callsite.callee) {
+            let mut args = args.into_iter();
+            let self_ = self.create_temp_if_necessary(args.next().unwrap(), callsite, caller_mir);
+            let tuple = self.create_temp_if_necessary(args.next().unwrap(), callsite, caller_mir);
+            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(Operand::Consume(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| Operand::Consume(self.create_temp_if_necessary(a, callsite, caller_mir)))
+                .collect()
+        }
+    }
+
+    /// If `arg` is already a temporary, returns it. Otherwise, introduces a fresh
+    /// temporary `T` and an instruction `T = arg`, and returns `T`.
+    fn create_temp_if_necessary(
+        &self,
+        arg: Operand<'tcx>,
+        callsite: &CallSite<'tcx>,
+        caller_mir: &mut Mir<'tcx>,
+    ) -> Lvalue<'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 {
-                if caller_mir.local_kind(local) == LocalKind::Temp {
-                    // Reuse the operand if it's a temporary already
-                    return a;
-                }
+
+        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 Lvalue::Local(local);
             }
+        }
 
-            debug!("Creating temp for argument");
-            // Otherwise, create a temporary for the arg
-            let arg = Rvalue::Use(a);
+        debug!("Creating temp for argument {:?}", arg);
+        // Otherwise, create a temporary for the arg
+        let arg = Rvalue::Use(arg);
 
-            let ty = arg.ty(caller_mir, tcx);
+        let ty = arg.ty(caller_mir, self.tcx);
 
-            let arg_tmp = LocalDecl::new_temp(ty, callsite.location.span);
-            let arg_tmp = caller_mir.local_decls.push(arg_tmp);
-            let arg_tmp = Lvalue::Local(arg_tmp);
+        let arg_tmp = LocalDecl::new_temp(ty, callsite.location.span);
+        let arg_tmp = caller_mir.local_decls.push(arg_tmp);
+        let arg_tmp = Lvalue::Local(arg_tmp);
 
-            let stmt = Statement {
-                source_info: callsite.location,
-                kind: StatementKind::Assign(arg_tmp.clone(), arg)
-            };
-            caller_mir[callsite.bb].statements.push(stmt);
-            Operand::Consume(arg_tmp)
-        }).collect()
+        let stmt = Statement {
+            source_info: callsite.location,
+            kind: StatementKind::Assign(arg_tmp.clone(), arg),
+        };
+        caller_mir[callsite.bb].statements.push(stmt);
+        arg_tmp
     }
 }
 
diff --git a/src/test/mir-opt/inline-closure.rs b/src/test/mir-opt/inline-closure.rs
new file mode 100644
index 00000000000..3f3428714d1
--- /dev/null
+++ b/src/test/mir-opt/inline-closure.rs
@@ -0,0 +1,42 @@
+// 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, q)
+}
+
+// END RUST SOURCE
+// START rustc.foo.Inline.after.mir
+// ...
+// bb0: {
+//     ...
+//     _3 = [closure@NodeId(28)];
+//     ...
+//     _4 = &_3;
+//     ...
+//     _6 = _2;
+//     ...
+//     _7 = _2;
+//     _5 = (_6, _7);
+//     _0 = (_5.0: i32);
+//     ...
+//     return;
+// }
+// ...
+// END rustc.foo.Inline.after.mir
\ No newline at end of file