[MIR-opt] MIR inlining handles closures incorrectly

[arielb1]

$ rustc -V
rustc 1.23.0-nightly (90ef3372e 2017-10-29)
$ rustc klosure.rs
$ ./klosure 
28
$ rustc klosure.rs -Z mir-opt-level=2
Function return type does not match operand type of return inst!
  ret { i32, [0 x i8], i32, [0 x i8] } %15
 i32LLVM ERROR: Broken function found, compilation aborted!

Code:

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)
}

Expected result: no LLVM assert, should work.

[arielb1]

So the problem is that inlining does not handle argument mapping correctly:

These MIRs should be inlined together correctly

// MIR for `foo`
// source = Fn(NodeId(7))
// pass_name = Inline
// disambiguator = before

fn foo(_1: T, _2: i32) -> i32 {
    let mut _0: i32;                     // return pointer
    scope 1 {
        let _3: [closure@klosure.rs:6:13: 6:24]; // "x" in scope 1 at klosure.rs:6:9: 6:10
    }
    scope 2 {
    }
    let mut _4: &[closure@klosure.rs:6:13: 6:24];
    let mut _5: (i32, i32);
    let mut _6: i32;
    let mut _7: i32;
    let mut _8: i32;
    let mut _9: i32;

    bb0: {                              
        StorageLive(_3);                 // scope 1 at klosure.rs:6:9: 6:10
        _3 = [closure@klosure.rs:6:13: 6:24]; // scope 1 at klosure.rs:6:13: 6:24
        StorageLive(_4);                 // scope 1 at klosure.rs:7:5: 7:6
        _4 = &_3;                        // scope 1 at klosure.rs:7:5: 7:6
        StorageLive(_5);                 // scope 1 at klosure.rs:7:5: 7:16
        StorageLive(_6);                 // scope 1 at klosure.rs:7:7: 7:10
        StorageLive(_7);                 // scope 1 at klosure.rs:7:7: 7:8
        _7 = _2;                         // scope 1 at klosure.rs:7:7: 7:8
        _6 = Mul(_7, const 2i32);        // scope 1 at klosure.rs:7:7: 7:10
        StorageDead(_7);                 // scope 1 at klosure.rs:7:10: 7:10
        StorageLive(_8);                 // scope 1 at klosure.rs:7:12: 7:15
        StorageLive(_9);                 // scope 1 at klosure.rs:7:12: 7:13
        _9 = _2;                         // scope 1 at klosure.rs:7:12: 7:13
        _8 = Mul(_9, const 3i32);        // scope 1 at klosure.rs:7:12: 7:15
        StorageDead(_9);                 // scope 1 at klosure.rs:7:15: 7:15
        _5 = (_6, _8);                   // scope 1 at klosure.rs:7:5: 7:16
        _0 = const std::ops::Fn::call(_4, _5) -> bb1; // scope 1 at klosure.rs:7:5: 7:16
    }

    bb1: {                              
        StorageDead(_5);                 // scope 1 at klosure.rs:7:16: 7:16
        StorageDead(_8);                 // scope 1 at klosure.rs:7:16: 7:16
        StorageDead(_6);                 // scope 1 at klosure.rs:7:16: 7:16
        StorageDead(_4);                 // scope 1 at klosure.rs:7:16: 7:16
        StorageDead(_3);                 // scope 0 at klosure.rs:8:2: 8:2
        return;                          // scope 0 at klosure.rs:8:2: 8:2
    }
}

// MIR for `foo::{{closure}}`
// source = Fn(NodeId(28))
// pass_name = Inline
// disambiguator = before

fn foo::{{closure}}(_1: &[closure@klosure.rs:6:13: 6:24], _2: i32, _3: i32) -> i32 {
    let mut _0: i32;                     // return pointer
    let mut _4: i32;

    bb0: {                              
        StorageLive(_4);                 // scope 0 at klosure.rs:6:22: 6:24
        _4 = _2;                         // scope 0 at klosure.rs:6:22: 6:24
        _0 = _4;                         // scope 0 at klosure.rs:6:22: 6:24
        StorageDead(_4);                 // scope 0 at klosure.rs:6:24: 6:24
        return;                          // scope 0 at klosure.rs:6:24: 6:24
    }
}

As you can see, the closure, which is called at the end of bb0, is passed its self-type and a tuple, but the closure signature itself takes 3 arguments - the self type and 2 parameters. When we inline, the inliner does not account for this difference, with this result:

// MIR for `foo`
// source = Fn(NodeId(7))
// pass_name = Inline
// disambiguator = after

fn foo(_1: T, _2: i32) -> i32 {
    let mut _0: i32;                     // return pointer
    scope 1 {
        let _3: [closure@klosure.rs:6:13: 6:24]; // "x" in scope 1 at klosure.rs:6:9: 6:10
        scope 3 {
        }
    }
    scope 2 {
    }
    let mut _4: &[closure@klosure.rs:6:13: 6:24];
    let mut _5: (i32, i32);
    let mut _6: i32;
    let mut _7: i32;
    let mut _8: i32;
    let mut _9: i32;

    bb0: {                              
        StorageLive(_3);                 // scope 1 at klosure.rs:6:9: 6:10
        _3 = [closure@klosure.rs:6:13: 6:24]; // scope 1 at klosure.rs:6:13: 6:24
        StorageLive(_4);                 // scope 1 at klosure.rs:7:5: 7:6
        _4 = &_3;                        // scope 1 at klosure.rs:7:5: 7:6
        StorageLive(_5);                 // scope 1 at klosure.rs:7:5: 7:16
        StorageLive(_6);                 // scope 1 at klosure.rs:7:7: 7:10
        StorageLive(_7);                 // scope 1 at klosure.rs:7:7: 7:8
        _7 = _2;                         // scope 1 at klosure.rs:7:7: 7:8
        _6 = Mul(_7, const 2i32);        // scope 1 at klosure.rs:7:7: 7:10
        StorageDead(_7);                 // scope 1 at klosure.rs:7:10: 7:10
        StorageLive(_8);                 // scope 1 at klosure.rs:7:12: 7:15
        StorageLive(_9);                 // scope 1 at klosure.rs:7:12: 7:13
        _9 = _2;                         // scope 1 at klosure.rs:7:12: 7:13
        _8 = Mul(_9, const 3i32);        // scope 1 at klosure.rs:7:12: 7:15
        StorageDead(_9);                 // scope 1 at klosure.rs:7:15: 7:15
        _5 = (_6, _8);                   // scope 1 at klosure.rs:7:5: 7:16
        // BUG: this should have been an access of _5.0 but instead accesses
        // the whole _5
        _0 = _5;                         // scope 3 at klosure.rs:6:22: 6:24
        StorageDead(_5);                 // scope 1 at klosure.rs:7:16: 7:16
        StorageDead(_8);                 // scope 1 at klosure.rs:7:16: 7:16
        StorageDead(_6);                 // scope 1 at klosure.rs:7:16: 7:16
        StorageDead(_4);                 // scope 1 at klosure.rs:7:16: 7:16
        StorageDead(_3);                 // scope 0 at klosure.rs:8:2: 8:2
        return;                          // scope 0 at klosure.rs:8:2: 8:2
    }
}

The inliner should be corrected to map accesses to the parameters of the closure to accesses to fields of the tupled parameters (later on, we might want an SROA-style pass to remove the tuple, but that's somewhat involved and not needed to fix this bug).

The inliner already maps parameters, with the code for that being in rustc_mir::transform::inline::Integrator:

rust/src/librustc_mir/transform/inline.rs

Lines 611 to 635 in 98e791e

	fn visit_local(&mut self,
	local: &mut Local,
	_ctxt: LvalueContext<'tcx>,
	_location: Location) {
	if *local == RETURN_POINTER {
	match self.destination {
	Lvalue::Local(l) => {
	*local = l;
	return;
	},
	ref lval => bug!("Return lvalue is {:?}, not local", lval)
	}
	}
	let idx = local.index() - 1;
	if idx < self.args.len() {
	match self.args[idx] {
	Operand::Consume(Lvalue::Local(l)) => {
	*local = l;
	return;
	},
	ref op => bug!("Arg operand `{:?}` is {:?}, not local", idx, op)
	}
	}
	*local = self.local_map[Local::new(idx - self.args.len())];
	}

That code needs to account for calls to closures, which can be identified by their disambiguated_data, as in the following:

rust/src/librustc_trans/common.rs

Lines 554 to 569 in 98e791e

	pub fn is_inline_instance<'a, 'tcx>(
	tcx: TyCtxt<'a, 'tcx, 'tcx>,
	instance: &ty::Instance<'tcx>
	) -> bool {
	let def_id = match instance.def {
	ty::InstanceDef::Item(def_id) => def_id,
	ty::InstanceDef::DropGlue(_, Some(_)) => return false,
	_ => return true
	};
	match tcx.def_key(def_id).disambiguated_data.data {
	DefPathData::StructCtor |
	DefPathData::EnumVariant(..) |
	DefPathData::ClosureExpr => true,
	_ => false
	}
	}

When encountering a closure, the code needs to be adapted to map locals to fields of the argument tuple.

If you need any help, feel free to ask me (@arielby), @nikomatsakis or @eddyb on the #rust-impl-period/WG-compiler-middle gitter channel - this would be a good way to get into MIR optimizations.