[mypyc] Merge generator and environment classes in simple cases

[JukkaL]

Mypyc used to always compile a generator or an async def into two classes: a generator and an environment. Now we combine these two classes in simple cases where it's clearly okay to do it (when there is no nesting). We could probably extend it to other use cases as well, including some nested functions, but this is a start.

This improves performance by reducing the number of instances that will be allocated. Also access to the environment object is slightly faster, though this is probably relatively minor. This helps calling async defs in particular, since they typically yield only a single value, so the objects are not used much until they are freed. Also generators that only yield a small number of values benefit from this.

The existing test cases provide decent test coverage. I previously added some additional tests in anticipation of this change.

This also reduces the amount of C code generated when compiling async defs and generators.

This speeds up this micro-benchmark on the order of 20%:

import asyncio
from time import time

async def inc(x: int) -> int:
    x = 1
    return x + 1


async def bench(n: int) -> int:
    x = 0
    for i in range(n):
        x = await inc(x)
    return x

asyncio.run(bench(1000))

t0 = time()
asyncio.run(bench(1000 * 1000 * 200))
print(time() - t0)

[JukkaL]

Here's IR for the inc async function from the PR description, before this PR (with some boilerplate removed):

def inc_gen.__mypyc_generator_helper__(__mypyc_self__, type, value, traceback, arg):
    __mypyc_self__ :: b.inc_gen
    type, value, traceback, arg :: object
    r0 :: b.inc_env
    r1 :: int
    r2 :: object
    r3 :: bit
    r4 :: bool
    r5, r6 :: int
    r7 :: object
    r8 :: bool
    r9 :: bit
    r10 :: bool
    r11 :: object
L0:
    r0 = __mypyc_self__.__mypyc_env__
    if is_error(r0) goto L13 (error at inc:1) else goto L1
L1:
    r1 = r0.__mypyc_next_label__
    if is_error(r1) goto L14 (error at inc:1) else goto L10
L2:
    r2 = load_address _Py_NoneStruct
    r3 = type != r2
    if r3 goto L15 else goto L5 :: bool
L3:
    CPyErr_SetObjectAndTraceback(type, value, traceback)
    if not 0 goto L13 (error at inc:1) else goto L4 :: bool
L4:
    unreachable
L5:
    r0.x = 2; r4 = is_error
    if not r4 goto L14 (error at inc:2) else goto L6 :: bool
L6:
    r5 = r0.x
    if is_error(r5) goto L14 (error at inc:3) else goto L7
L7:
    r6 = CPyTagged_Add(r5, 2)
    dec_ref r5 :: int
    r7 = box(int, r6)
    r0.__mypyc_next_label__ = -2; r8 = is_error
    dec_ref r0
    if not r8 goto L16 (error at inc:3) else goto L8 :: bool
L8:
    CPyGen_SetStopIterationValue(r7)
    dec_ref r7
    if not 0 goto L13 else goto L9 :: bool
L9:
    unreachable
L10:
    r9 = r1 == 0
    dec_ref r1 :: int
    if r9 goto L2 else goto L17 :: bool
L11:
    r10 = raise StopIteration
    if not r10 goto L13 (error at inc:1) else goto L12 :: bool
L12:
    unreachable
L13:
    r11 = <error> :: object
    return r11
L14:
    dec_ref r0
    goto L13
L15:
    dec_ref r0
    goto L3
L16:
    dec_ref r7
    goto L13
L17:
    dec_ref r0
    goto L11

def inc_gen.__next__(__mypyc_self__):
    __mypyc_self__ :: b.inc_gen
    r0, r1, r2 :: object
L0:
    r0 = load_address _Py_NoneStruct
    r1 = inc_gen.__mypyc_generator_helper__(__mypyc_self__, r0, r0, r0, r0)
    if is_error(r1) goto L2 else goto L1
L1:
    return r1
L2:
    r2 = <error> :: object
    return r2

def inc_gen.__iter__(__mypyc_self__):
    __mypyc_self__ :: b.inc_gen
L0:
    inc_ref __mypyc_self__
    return __mypyc_self__

def inc_gen.__await__(__mypyc_self__):
    __mypyc_self__ :: b.inc_gen
L0:
    inc_ref __mypyc_self__
    return __mypyc_self__

def inc(x):
    x :: int
    r0 :: b.inc_env
    r1 :: bool
    r2 :: b.inc_gen
    r3, r4 :: bool
    r5 :: object
L0:
    r0 = inc_env()
    if is_error(r0) goto L6 (error at inc:1) else goto L1
L1:
    inc_ref x :: int
    r0.x = x; r1 = is_error
    if not r1 goto L7 (error at inc:1) else goto L2 :: bool
L2:
    r2 = inc_gen()
    if is_error(r2) goto L7 (error at inc:1) else goto L3
L3:
    inc_ref r0
    r2.__mypyc_env__ = r0; r3 = is_error
    if not r3 goto L8 (error at inc:1) else goto L4 :: bool
L4:
    r0.__mypyc_next_label__ = 0; r4 = is_error
    dec_ref r0
    if not r4 goto L9 (error at inc:1) else goto L5 :: bool
L5:
    return r2
L6:
    r5 = <error> :: object
    return r5
L7:
    dec_ref r0
    goto L6
L8:
    dec_ref r0
    dec_ref r2
    goto L6
L9:
    dec_ref r2
    goto L6

Here's the IR after this PR, when merging is enabled:

def inc_gen.__mypyc_generator_helper__(__mypyc_self__, type, value, traceback, arg):
    __mypyc_self__ :: b.inc_gen
    type, value, traceback, arg :: object
    r0 :: int
    r1 :: object
    r2 :: bit
    r3 :: bool
    r4, r5 :: int
    r6 :: object
    r7 :: bool
    r8 :: bit
    r9 :: bool
    r10 :: object
L0:
    r0 = __mypyc_self__.__mypyc_next_label__
    if is_error(r0) goto L12 (error at inc:1) else goto L9
L1:
    r1 = load_address _Py_NoneStruct
    r2 = type != r1
    if r2 goto L2 else goto L4 :: bool
L2:
    CPyErr_SetObjectAndTraceback(type, value, traceback)
    if not 0 goto L12 (error at inc:1) else goto L3 :: bool
L3:
    unreachable
L4:
    __mypyc_self__.x = 2; r3 = is_error
    if not r3 goto L12 (error at inc:2) else goto L5 :: bool
L5:
    r4 = __mypyc_self__.x
    if is_error(r4) goto L12 (error at inc:3) else goto L6
L6:
    r5 = CPyTagged_Add(r4, 2)
    dec_ref r4 :: int
    r6 = box(int, r5)
    __mypyc_self__.__mypyc_next_label__ = -2; r7 = is_error
    if not r7 goto L13 (error at inc:3) else goto L7 :: bool
L7:
    CPyGen_SetStopIterationValue(r6)
    dec_ref r6
    if not 0 goto L12 else goto L8 :: bool
L8:
    unreachable
L9:
    r8 = r0 == 0
    dec_ref r0 :: int
    if r8 goto L1 else goto L10 :: bool
L10:
    r9 = raise StopIteration
    if not r9 goto L12 (error at inc:1) else goto L11 :: bool
L11:
    unreachable
L12:
    r10 = <error> :: object
    return r10
L13:
    dec_ref r6
    goto L12

def inc_gen.__next__(__mypyc_self__):
    __mypyc_self__ :: b.inc_gen
    r0, r1, r2 :: object
L0:
    r0 = load_address _Py_NoneStruct
    r1 = inc_gen.__mypyc_generator_helper__(__mypyc_self__, r0, r0, r0, r0)
    if is_error(r1) goto L2 else goto L1
L1:
    return r1
L2:
    r2 = <error> :: object
    return r2

def inc_gen.__iter__(__mypyc_self__):
    __mypyc_self__ :: b.inc_gen
L0:
    inc_ref __mypyc_self__
    return __mypyc_self__

def inc_gen.__await__(__mypyc_self__):
    __mypyc_self__ :: b.inc_gen
L0:
    inc_ref __mypyc_self__
    return __mypyc_self__

def inc(x):
    x :: int
    r0 :: b.inc_gen
    r1, r2 :: bool
    r3 :: object
L0:
    r0 = inc_gen()
    if is_error(r0) goto L4 (error at inc:1) else goto L1
L1:
    r0.__mypyc_next_label__ = 0; r1 = is_error
    if not r1 goto L5 (error at inc:1) else goto L2 :: bool
L2:
    inc_ref x :: int
    r0.x = x; r2 = is_error
    if not r2 goto L5 (error at inc:1) else goto L3 :: bool
L3:
    return r0
L4:
    r3 = <error> :: object
    return r3
L5:
    dec_ref r0
    goto L4