Generator size: unwinding and drops force extra generator state allocation

[Nemo157]

#![feature(generators, generator_trait)]

use std::ops::Generator;

struct Foo([u8; 1024]);

impl Drop for Foo {
    fn drop(&mut self) {}
}

fn simple() -> impl Generator<Yield = (), Return = ()> {
    static || {
        let first = Foo([0; 1024]);
        let _second = first;
        yield;
    }
}

fn complex() -> impl Generator<Yield = (), Return = ()> {
    static || {
        let first = Foo([0; 1024]);
        { foo(); fn foo() {} }
        let _second = first;
        yield;
    }
}

fn main() {
    dbg!(std::mem::size_of_val(&simple()));
    dbg!(std::mem::size_of_val(&complex()));
}

The two generators returned by simple and complex should be equivalent, but complex takes twice as much space:

[foo.rs:29] std::mem::size_of_val(&simple()) = 1028
[foo.rs:30] std::mem::size_of_val(&complex()) = 2056

---

Dumping out the MIR (with rustc 1.34.0-nightly (f66e4697a 2019-02-20)) shows an issue with how unwinding from foo interacts with the two stack slots for first and _second, using a dynamic drop flag means that first is "live" through the path that goes through the yield, even though the drop flag is guaranteed to be false. (The below graph shows the basic blocks, with the psuedo-code run in them and which variables are alive when exiting the block):

[Nemo157]

@rustbot modify labels: A-generators and T-compiler.

[matprec]

Is this related to #52924?

[Nemo157]

@MSleepyPanda in as much as it’s about generators being too big. The specific optimisation proposed there won’t help here as first and second are both live over the same yield point. What should really happen is that first is not kept live across the yield at all and it should be allocated in the resume function stack instead of the generator state. (And then some sort of copy-elision optimisation might eliminate that allocation and use the allocation in the generator state directly, but IMO that’s less important (and probably more difficult) than ensuring the memory usage is reduced).

[cramertj]

cc @tmandry

[andreytkachenko]

Another example, without drop:

#![feature(generators, generator_trait)]

use std::ops::Generator;

struct Foo([u8; 1024]);

fn simple() -> impl Generator<Yield = (), Return = ()> {
    static || {
        let first = Foo([0; 1024]);
        let _second = first;
        yield;
    }
}

fn complex() -> impl Generator<Yield = (), Return = ()> {
    static || {
        fn foo(_: &mut Foo) {}
        
        let mut first = Foo([0; 1024]);
        foo(&mut first);
        let mut second = first;
        foo(&mut second);
        let mut third = second;
        foo(&mut third);
        let mut _fourth = third;
        
        yield;
    }
}

fn main() {
    dbg!(std::mem::size_of_val(&simple()));
    dbg!(std::mem::size_of_val(&complex()));
}

outputs

[src/main.rs:32] std::mem::size_of_val(&simple()) = 4
[src/main.rs:33] std::mem::size_of_val(&complex()) = 3076

[tmandry]

I was thinking, we can solve this by adding the following rules to our MaybeStorageLive dataflow analysis (possibly being renamed to RequiresStorage):

1. (Existing) StorageLive(x) => mark x live
2. (Existing) StorageDead(x) => mark x dead
3. If a local is moved from, and has never been mutably borrowed, mark it dead
4. If (any part of) a local is initialized, mark it live

We must not optimize away storage of locals that are mutably borrowed, because as @matthewjasper notes in #61430, it isn't decided that the following is UB:

let mut x = String::new();
let p = &mut x as *mut String;
let y = x;
p.write(String::new());

It's an open question of whether we can say "the local hasn't been mutably borrowed up to here" when evaluating rule 3. I'd prefer to make the optimization as smart as we can, but MIR probably allows borrowing a moved-from value and mutating it.

Is this sound?

cc @cramertj @eddyb @matthewjasper @RalfJung @Zoxc

[eddyb]

The "mutably" of "mutably borrowed" is a red herring IMO, unless you want to check for Freeze, which will conservatively default to "may contain interior mutability" once generic parameters are thrown into the mix.

[RalfJung]

@tmandry Interesting. How bad would it be to relax this to "if a local is moved from and never has had its address taken"? Then we can be sure without any assumptions about Stacked Borrows that direct accesses to the local variable are the only way to observe it, and those will be prevented after a move. This would also alleviate @eddyb's concern I think.

Also, what is the granularity here? Without Stacked Borrows assumptions we can only do this on a per-local level, not on a per-field-of-a-struct level. Taking the address of one field leaks the address of other fields (if layout assumptions are made).