Generator size: borrowed variables are assumed live across following yield points

[Matthias247]

Maybe a duplicate of #52924, but maybe also something else.

I observed that the sizes of Futures generated by async fns can grow exponentially.
The following code shows an async fn, which produces a 1kB future. Each layering in another async fn doubles it's size:

#![feature(async_await)]

async fn i_am_1kb() -> bool
{
    let x: [u8; 1*1024] = [0; 1*1024];
    async{}.await;
    let _sum: u8 = x.iter().sum();
    true
}

fn main() {
    let fut1 = i_am_1kb();
    dbg!(std::mem::size_of_val(&fut1));

    let composed_1 = async {
        let inner = i_am_1kb();
        inner.await;
    };
    dbg!(std::mem::size_of_val(&composed_1));

    let composed_2 = async {
        let inner = i_am_1kb();
        dbg!(std::mem::size_of_val(&inner));
        inner.await;
    };
    dbg!(std::mem::size_of_val(&composed_2));

    let composed_3 = async {
        let inner = async {
            let inner = async {
                i_am_1kb().await;
            };
            dbg!(std::mem::size_of_val(&inner));
            inner.await;
        };
        dbg!(std::mem::size_of_val(&inner));
        inner.await;
    };
    dbg!(std::mem::size_of_val(&composed_3));
}

Output:

[src/main.rs:16] std::mem::size_of_val(&fut1) = 1032
[src/main.rs:22] std::mem::size_of_val(&composed_1) = 1036
[src/main.rs:29] std::mem::size_of_val(&composed_2) = 2072
[src/main.rs:44] std::mem::size_of_val(&composed_3) = 4168

It doesn't matter whether the statement between the future generation and await! references the future or not. A simply println("") will have the same effect.
Only if the future is directly awaited (as in composed_1) the size will stay constant.

cc @cramertj , @nikomatsakis , @Nemo157

[cramertj]

Closing as a sub-issue of #52924. (I've edited the top message in that thread to reference this one)

[cramertj]

cc @tmandry @Zoxc

[Nemo157]

There definitely seems to be something causing locals to be unnecessarily put into the generator struct instead of staying as true-locals. Using a simplified function (with the same i_am_1kb as above)

async fn composed() {
    let inner = i_am_1kb();
    { let _foo = &inner; }
    await!(inner);
}

and running cargo rustc -- -Z dump-mir=generator to dump the mir, the liveness analysis shows that inner is correctly considered dead after being moved to pinned in the await! macro (and so is not alive over a yield), but it is still being put in the generator.

[Zoxc]

@Nemo157 The generator transformation conservatively assumes that any borrow can be converted to a raw pointer and the locals can be accessed with that until their storage slot is dead. That's why inner is considered live during the await! here.

[cramertj]

Yeah, @eddyb and I discussed this at the all-hands, and that making a type implement Copy is actually a potential perf regression here, which is weird. For non-Copy types, you can assume no accesses after moving out of them, but for Copy types you can't necessarily do this.

[Nemo157]

Ok, so an optimisation to fix this example would be to add a less conservative check that can see when the borrow definitely wasn’t converted to a raw pointer. That seems relatively straightforward to check when the borrow never enters any unsafe code.

[cramertj]

add a less conservative check that can see when the borrow definitely wasn’t converted to a raw pointer

Note that for this to be very useful at all it would have to be able to see through functions (via MIR inlining) and intrinsics (e.g. size_of_val above).

[Nemo157]

Can it not trust the lifetimes on safe functions signatures? size_of_val does not have a lifetime dependency so should be relied on to not stash a raw pointer to the reference away somewhere. (I guess this is an UCG question whether safe function boundaries are barriers that require safety to be upheld, and whether future unsafe code can allow prior safe code to violate lifetimes, i.e. is something like this sound or not).

[cramertj]

No, lifetimes in function signatures cannot necessarily be used to determine the scope of accesses to the resulting pointer. Without a memory model it's completely unclear when accesses would or wouldn't be allowed to the underlying memory. @RalfJung's work on stacked borrows is the only thing I'm aware of that would allow proper analysis, and in general anywhere there's a ref-to-ptr conversion, all bets are sort of off.