Skip to content

Commit

Permalink
fix closure inlining by spilling arguments to a temporary
Browse files Browse the repository at this point in the history
  • Loading branch information
nikomatsakis committed Nov 19, 2017
1 parent 83f5a96 commit df6fdbc
Show file tree
Hide file tree
Showing 3 changed files with 93 additions and 10 deletions.
47 changes: 39 additions & 8 deletions src/librustc_mir/transform/inline.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ use rustc::ty::{self, Instance, Ty, TyCtxt, TypeFoldable};
use rustc::ty::subst::{Subst,Substs};

use std::collections::VecDeque;
use std::iter;
use transform::{MirPass, MirSource};
use super::simplify::{remove_dead_blocks, CfgSimplifier};

Expand Down Expand Up @@ -558,8 +559,29 @@ impl<'a, 'tcx> Inliner<'a, '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.
// There is a bit of a mismatch between the *caller* of a closure and the *callee*.
// The caller provides the arguments wrapped up in a tuple:
//
// tuple_tmp = (a, b, c)
// Fn::call(closure_ref, tuple_tmp)
//
// meanwhile the closure body expects the arguments (here, `a`, `b`, and `c`)
// as distinct arguments. (This is the "rust-call" ABI hack.) Normally, trans has
// the job of unpacking this tuple. But here, we are trans. =) So we want to create
// a vector like
//
// [closure_ref, tuple_tmp.0, tuple_tmp.1, tuple_tmp.2]
//
// Except for one tiny wrinkle: we don't actually want `tuple_tmp.0`. It's more convenient
// if we "spill" that into *another* temporary, so that we can map the argument
// variable in the callee MIR directly to an argument variable on our side.
// So we introduce temporaries like:
//
// tmp0 = tuple_tmp.0
// tmp1 = tuple_tmp.1
// tmp2 = tuple_tmp.2
//
// and the vector is `[closure_ref, tmp0, tmp1, tmp2]`.
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);
Expand All @@ -572,12 +594,21 @@ impl<'a, 'tcx> Inliner<'a, 'tcx> {
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
// The `closure_ref` in our example above.
let closure_ref_arg = iter::once(Operand::Consume(self_));

// The `tmp0`, `tmp1`, and `tmp2` in our example abonve.
let tuple_tmp_args =
tuple_tys.iter().enumerate().map(|(i, ty)| {
// This is e.g. `tuple_tmp.0` in our example above.
let tuple_field = Operand::Consume(tuple.clone().field(Field::new(i), ty));

// Spill to a local to make e.g. `tmp0`.
let tmp = self.create_temp_if_necessary(tuple_field, callsite, caller_mir);
Operand::Consume(tmp)
});

closure_ref_arg.chain(tuple_tmp_args).collect()
} else {
args.into_iter()
.map(|a| Operand::Consume(self.create_temp_if_necessary(a, callsite, caller_mir)))
Expand Down
50 changes: 50 additions & 0 deletions src/test/mir-opt/inline-closure-borrows-arg.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
// 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,
// even when (#45894)

fn main() {
println!("{}", foo(0, &14));
}

fn foo<T: Copy>(_t: T, q: &i32) -> i32 {
let x = |r: &i32, _s: &i32| {
let variable = &*r;
*variable
};
x(q, q)
}

// END RUST SOURCE
// START rustc.foo.Inline.after.mir
// ...
// bb0: {
// ...
// _3 = [closure@NodeId(39)];
// ...
// _4 = &_3;
// ...
// _6 = &(*_2);
// ...
// _7 = &(*_2);
// _5 = (_6, _7);
// _9 = (_5.0: &i32);
// _10 = (_5.1: &i32);
// StorageLive(_8);
// _8 = (*_9);
// _0 = _8;
// ...
// return;
// }
// ...
// END rustc.foo.Inline.after.mir
6 changes: 4 additions & 2 deletions src/test/mir-opt/inline-closure.rs
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,11 @@ fn foo<T: Copy>(_t: T, q: i32) -> i32 {
// ...
// _7 = _2;
// _5 = (_6, _7);
// _0 = (_5.0: i32);
// _8 = (_5.0: i32);
// _9 = (_5.1: i32);
// _0 = _8;
// ...
// return;
// }
// ...
// END rustc.foo.Inline.after.mir
// END rustc.foo.Inline.after.mir

0 comments on commit df6fdbc

Please sign in to comment.