Stack space in async closures is not shared between branches in debug builds

[khuey]

Consider:

#![feature(test)]
use std::future::Future;
use std::hint::black_box;
use std::pin::Pin;

#[derive(PartialEq, Eq)]
enum Kind {
    Thing1,
    Thing2,
    Thing3,
}

impl Into<i32> for Kind {
    fn into(self) -> i32 {
        match self {
            Kind::Thing1 => 42,
            Kind::Thing2 => -1,
            Kind::Thing3 => 0,
        }
    }
}

fn f1(kind: Kind) -> Pin<Box<dyn Future<Output = i32>>> {
    Box::pin(async move {
        match kind {
            Kind::Thing1 => return kind.into(),
            Kind::Thing2 => return kind.into(),
            Kind::Thing3 => return kind.into(),
        }
    })
}

fn f2(kind: Kind) -> Pin<Box<dyn Future<Output = i32>>> {
    match kind {
        Kind::Thing1 => Box::pin(async move { kind.into() }),
        Kind::Thing2 => Box::pin(async move { kind.into() }),
        Kind::Thing3 => Box::pin(async move { kind.into() }),
    }
}

fn f3(kind: Kind) -> Pin<Box<dyn Future<Output = i32>>> {
    Box::pin(async move {
        if kind == Kind::Thing1 {
            return kind.into();
        } else if kind == Kind::Thing2 {
            return kind.into();
        } else {
            return kind.into();
        }
    })
}

fn main() {
    let f_1 = f1(Kind::Thing2);
    let f_2 = f2(Kind::Thing2);
    let f_3 = f3(Kind::Thing2);
    black_box(Kind::Thing1);
    black_box(Kind::Thing3);
    black_box(f_1);
    black_box(f_2);
    black_box(f_3);
}

If you build this (on 2020-05-14 a74d186) you'll see that the stack frames for f1::{{closure}} and f3::{{closure}} are bigger than those for the f2::{{closure}s. If you try to add or remove a branch you'll also see that the size of f1::{{closure}} and f3::{{closure}} stacks vary accordingly.

This is frustrating because it means that branch heavy code blows out the stack much faster in a debug build than it does in an opt build when recursing. I saw an order of magnitude difference in stack size in one of my functions between debug and opt because of this.

[jonas-schievink]

Is this really limited to async code?

[khuey]

I don't know. There's not really any equivalent to f2 in non-async code, so it's not as easy to measure.

[khuey]

From a quick test on a non-async function with a bunch of trivial branches the stack space is constant regardless of how many branches I add.

[khuey]

I believe the problem here is panic handling. In create_generator_resume_function the generator MIR code takes every basic block in the generator with a terminator that a) can unwind but b) doesn't have an unwind handler already and c) isn't in a cleanup block and points its unwind handler at the generator's poisoned block.

This means any MIR of the form

    bbN: {
        StorageLive(_X);
        _Y = call of some sort -> [return: bbN+1];
    }

    bbN+1: {
        StorageDead(_X);
        goto -> bbN+2;
    }

becomes

    bbN: {
        StorageLive(_X);
        _Y = call of some sort -> [return: bbN+1, unwind: bbPoisoned];
    }

    bbN+1: {
        StorageDead(_X);
        goto -> bbN+2;
    }

    bbPoisoned (cleanup): {
        discriminant(Generator) = 2;
        resume;
    }

So now every storage location that was live across something that could theoretically panic is now live at the same time on the poisoned block, which means none of them can share stack space. Thus the size of the stack blows out enormously. This is also consistent with my observation that adding a match branch that calls no functions but merely returns a constant does not increase the stack size.

So it seems to me that, in essence, #61373 has been regressed here. Maybe @tmandry has some idea what's going on?

[khuey]

And, to be clear, I simplified this example from some async code with a 100KB stack frame, so this is (at least a part of) a rather serious problem.

[jonas-schievink]

All generators in the posted code have a state size of 2 fwiw

[khuey]

Sorry if I wasn't clear: the values here are not live across yield points so they don't take up space in the generator itself, just on the C stack.

[jonas-schievink]

Yeah, you're right. It's just that #61373 was meant to reduce the state size, not the stack usage. This issue was likely introduced by #69814 then (even though it did not affect state size negatively).